TensorDict: A GPU-accelerated Python dictionary

If you're tired of dealing with boilerplate code when working with multiple tensors in PyTorch, or struggle with scaling operations to large collections of tensors, look no further! TensorDict is a PyTorch primitive that makes it easy to build dictionaries of tensors (and non-tensors) with a focus on composability and performance.

With TensorDict, you can apply the same operation to a collection of tensors with ease, eliminating the need for tedious loops and improving your code's readability. Our API is designed to be intuitive, mirroring the torch.Tensor API, so you can reshape, split, concatenate, clone, and more with familiar syntax.

But that's not all - TensorDict also optimizes operations under the hood, resulting in significant speedups. For example, casting a large collection of tensors to GPU can be up to 2x faster than using a regular Python loop. We used fused CUDA kernels for arithmetic ops so you can code an optimizer like ADAM in 5 lines of code (like you would do for a single tensor) and it will run much faster in eager and compile modes than you would have done with a regular for loop over your tensors.

Plus, our library comes with a GPU-accelerated dataclass (@tensorclass) for even more performance gains.

Other key features include: - torch. Compile (PT2) compatibility - Consolidation into a single storage for fast node-to-node communication - Support for memory-mapping and shared memory - TensorDict can be used as lightweight substitute to a dataloader thanks to `TensorDict.map_iter`, and do preproc on-device, outperforming regular dataloading speed by orders of magnitude (check the tutorials!)

We're excited to share TensorDict with the community and invite your feedback and contributions!

Try it out and let us know what you think!

a very convenient tool for any machine learning project in PyTorch! Speaking from experience.