TensorDict
TensorDict is a batched, nested dict[str, Tensor] that behaves like a tensor.
Move it, slice it, reshape it, stack it, save it, compile it, or do arithmetic on
it: every tensor leaf follows the same operation, and one shared batch_size
keeps the structure honest.
TensorDict(batch_size=[32])
|-- obs: Tensor[32, 128]
|-- action: Tensor[32]
|-- reward: Tensor[32]
`-- next:
`-- obs: Tensor[32, 128]
30-second demo | Why TensorDict | What is new in 0.13 | Patterns | Installation | Ecosystem | Citation
30-second demo
import torch
from tensordict import TensorDict
batch = TensorDict(
{
"obs": torch.randn(32, 128),
"action": torch.randint(0, 4, (32,)),
"reward": torch.randn(32),
"next": {"obs": torch.randn(32, 128)},
},
batch_size=[32],
)
mini = batch[:8] # slices every leaf
device = "cuda" if torch.cuda.is_available() else "cpu"
on_device = batch.to(device) # moves every leaf; non-blocking internally
scaled = batch * 0.5 # arithmetic on the whole structure
merged = batch + batch # leaf-wise TensorDict arithmetic
stacked = torch.stack([batch, batch], 0)
print(mini.shape) # torch.Size([8])
print(stacked.shape) # torch.Size([2, 32])
The object remains a mapping, but the batch acts like a tensor. That is the point: write the operation once, apply it to every tensor that belongs to the same example, rollout, batch, parameter set, or dataset shard.
Why TensorDict
Plain dictionaries are flexible. TensorDict keeps that flexibility and adds the parts tensor programs need once the code gets serious.
With a plain dict |
With TensorDict |
|---|---|
| Manually keep leading dimensions aligned | One batch_size validates the structure |
Repeat .to(device) for every tensor |
td.to(device) moves the full batch |
| Hand-roll slicing, stacking, reshaping | td[:32], torch.stack, td.reshape |
| Manually recurse through nested state | Nested keys are first-class |
| Duplicate arithmetic over leaves | td + td, td * scalar, td.abs() |
| Invent checkpoint formats | td.save, td.memmap, load_memmap |
| Hope generic code keeps working | PyTorch-native APIs, torch.compile coverage |
Use TensorDict when the unit of data is not one tensor anymore, but it should still move through your program like one tensor.
Performance is part of the API
TensorDict is not just syntax for recursive Python loops. Core paths are built for high-throughput PyTorch workloads:
- Arithmetic dispatch: operations such as
td + td,td * 0.5,td.abs()and in-place variants apply directly to leaves and use PyTorch foreach kernels where available. - Device and host transfers: D2H and H2D copies are dispatched across the
full structure. TensorDict uses non-blocking leaf transfers internally when
possible, so the common path is just
td.to(device); passnon_blocking=Falseonly when you need an explicitly synchronous transfer. - Shape operations without boilerplate: indexing,
view,reshape,permute,unsqueeze,squeeze,flatten,unflatten,stackandcatoperate on the batch structure rather than on hand-maintained lists of leaves. - Low-allocation workflows: lazy stacks, preallocation, memory mapping and
inplace=Trueshape-changing operations help reduce peak memory in data-heavy pipelines. - Compile-aware internals: TensorDict is used in compiled training and RL
loops, and the codebase carries dedicated
torch.compilecoverage for hot paths.
For deeper numbers, see the benchmark notes.
What is new in 0.13
TensorDict 0.13 focuses on making structured tensor programs more practical in large training systems:
- Tabular import/export for pandas, CSV, Parquet and JSON workflows.
- More
inplace=Trueshape operations, includinggather,repeat,repeat_interleave,roll,reshape,flatten,unflattenandcontiguous. - Improved
torch.compilebehavior for TensorClass initialization, dynamic-shape export, locking paths and shallow clones. - Safer memmap filenames by default through robust key encoding.
- A migration path for module state preservation with
to_module(..., preserve_module_state=...). - CPU-only release wheels for TensorDict, avoiding duplicate GPU wheel artifacts for a package whose compiled extension is device-independent.
Patterns
One batch through the whole training step
TensorDict lets datasets, models and losses agree on one container instead of a long argument list.
for batch in dataloader:
batch = batch.to(device)
batch = model(batch)
loss = loss_module(batch)
loss.backward()
optimizer.step()
optimizer.zero_grad()
That loop can stay stable while the schema changes from classification to segmentation, RL rollouts, model-based prediction or LLM post-training batches.
Nested data without custom plumbing
td = TensorDict(
{
"agents": {
"policy": torch.randn(64, 8),
"value": torch.randn(64, 1),
},
"env": {
"reward": torch.randn(64),
"done": torch.zeros(64, dtype=torch.bool),
},
},
batch_size=[64],
)
policy = td["agents", "policy"]
td["env", "reward"] = td["env", "reward"].clip(-1, 1)
Nested keys are part of the API, not an afterthought.
Functional modules and parameter sets
TensorDict can hold module parameters, swap them into modules, vectorize over ensembles and make model state explicit.
from tensordict import TensorDict
params = TensorDict.from_module(module)
with params.to_module(module, preserve_module_state=True):
out = module(inputs)
This is the same foundation used by TorchRL modules and functional training utilities.
Checkpoint and share large tensor batches
td = TensorDict({"tokens": tokens, "scores": scores}, batch_size=[n])
td.memmap("/path/to/private/batch") # memory-map every leaf
reloaded = TensorDict.load_memmap("/path/to/private/batch")
Memory-mapped TensorDicts are useful for large offline datasets, replay buffers, inter-process handoff and checkpointed intermediate state.
Key features
- Tensor-like collection ops: indexing, slicing, device casting, dtype casting, reshaping, stacking and concatenation. [tutorial]
- Nested structures with tuple keys and predictable batch semantics. [tutorial]
- Fast memory workflows: asynchronous transfers, memmap, consolidated tensors, lazy stacks and preallocation. [tutorial]
- Functional programming with parameter TensorDicts,
to_moduleand compatibility withtorch.vmap. [tutorial] @tensorclass: a tensor-aware dataclass for structured tensor objects. [tutorial]- Distributed and multiprocessed pipelines across workers, devices and machines. [doc]
- Serialization and memory mapping for efficient checkpointing and dataset storage. [doc]
For a longer tour, start with GETTING_STARTED.md or the online documentation.
Installation
With pip:
pip install tensordict
With conda:
conda install -c conda-forge tensordict
Nightly builds:
pip install tensordict-nightly
From source with an existing PyTorch install:
pip install -e . --no-deps
If you use uv with PyTorch nightlies, keep torch pinned to the PyTorch wheel
index or install TensorDict with --no-deps so the resolver does not replace
your existing PyTorch build:
uv pip install -e . --no-deps
uv pip install -e . --prerelease=allow -f "https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html"
Ecosystem
TensorDict started in reinforcement learning, where batches quickly become nested trajectories. It is now used anywhere tensor batches are structured data: RL rollouts, LLM post-training samples, robotics trajectories, simulation state, model parameters, checkpointed datasets and scientific pipelines.
| Domain | Projects |
|---|---|
| Reinforcement Learning | TorchRL (PyTorch), DreamerV3-torch, Dreamer4, SkyRL |
| LLM Post-Training | verl, ROLL (Alibaba), LMFlow, LoongFlow (Baidu) |
| Robotics and Simulation | MuJoCo Playground (Google DeepMind), ProtoMotions (NVIDIA), holosoma (Amazon) |
| Physics and Scientific ML | PhysicsNeMo (NVIDIA) |
| Genomics | Medaka (Oxford Nanopore) |
Citation
If you use TensorDict, please cite the TorchRL paper:
@misc{bou2023torchrl,
title={TorchRL: A data-driven decision-making library for PyTorch},
author={Albert Bou and Matteo Bettini and Sebastian Dittert and Vikash Kumar and Shagun Sodhani and Xiaomeng Yang and Gianni De Fabritiis and Vincent Moens},
year={2023},
eprint={2306.00577},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
License
TensorDict is licensed under the MIT License. See LICENSE for details.
1 maintainer
Release files for tensordict-nightly 2026.9.24
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Built distributions (wheels)
Total release size: 6.2 MB
Release files / tensordict_nightly-2026.9.24-cp314-cp314t-manylinux_2_28_aarch64.whl
| Download URL | tensordict_nightly-2026.9.24-cp314-cp314t-manylinux_2_28_aarch64.whl |
|---|---|
| Size | 595.2 kB |
| Tags | CPython 3.14 CPython 3.14 free-threading Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
7d69e732a62435cbaebd1d2d4b19b13a13d34e7b93d9c6c5feb10a855b023c98
|
|
BLAKE2b-256 checksum How to use checksums |
1c0274f67beb715392ca68ed7f32ab83e8f649c93c7f121db78925ec555b7a4f
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.10.12
|
Release files / tensordict_nightly-2026.9.24-cp314-cp314-win_amd64.whl
| Download URL | tensordict_nightly-2026.9.24-cp314-cp314-win_amd64.whl |
|---|---|
| Size | 656.1 kB |
| Tags | CPython 3.14 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
4a4f1299aa4d67e1718caa7073a327a75c8cf8f2341e8366d5fe4e12c3daea0f
|
|
BLAKE2b-256 checksum How to use checksums |
25173a3f46c4250cce2c5f543382ad95227f6df85aa22bc8df11e39be31cd01c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.10
|
Release files / tensordict_nightly-2026.9.24-cp314-cp314-manylinux_2_28_aarch64.whl
| Download URL | tensordict_nightly-2026.9.24-cp314-cp314-manylinux_2_28_aarch64.whl |
|---|---|
| Size | 594.0 kB |
| Tags | CPython 3.14 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
ec64907b85fbdf891211396faec71562921b244a2082ed730ac22096a72fb07a
|
|
BLAKE2b-256 checksum How to use checksums |
1e57f3a1eb0908febbc954669fca53ea1b5024f21f98841870364cad123ec02c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.10.12
|
Release files / tensordict_nightly-2026.9.24-cp313-cp313-win_amd64.whl
| Download URL | tensordict_nightly-2026.9.24-cp313-cp313-win_amd64.whl |
|---|---|
| Size | 654.1 kB |
| Tags | CPython 3.13 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
dc8f0c12af1659fce72b79d28bdfed37d541f006b6348e83ca3e24f9900d5dd5
|
|
BLAKE2b-256 checksum How to use checksums |
48b4f42a93a597b3c9003c2bb72ac13050498a17d3108aa2a1f10dc6d288bcc9
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.10
|
Release files / tensordict_nightly-2026.9.24-cp313-cp313-manylinux_2_28_aarch64.whl
| Download URL | tensordict_nightly-2026.9.24-cp313-cp313-manylinux_2_28_aarch64.whl |
|---|---|
| Size | 593.2 kB |
| Tags | CPython 3.13 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
e4b3615746c65780c1f76a9a645f8a2dd6afc827dcce2fd5f53fec63ee7ced24
|
|
BLAKE2b-256 checksum How to use checksums |
ecefc42c2908a3391808d0d381249331e626f80467e1a40e50d2e81a0572a453
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.10.12
|
Release files / tensordict_nightly-2026.9.24-cp312-cp312-win_amd64.whl
| Download URL | tensordict_nightly-2026.9.24-cp312-cp312-win_amd64.whl |
|---|---|
| Size | 654.0 kB |
| Tags | CPython 3.12 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
7945c1757cdee198a1ba125d5560a65135773b048a3828ef4199f12dd3b68be2
|
|
BLAKE2b-256 checksum How to use checksums |
7ff82f0b54f6c6c21d88e2d1619b6b2f882fc9a53edb217ca8b88e0074abcb6e
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.10
|
Release files / tensordict_nightly-2026.9.24-cp312-cp312-manylinux_2_28_aarch64.whl
| Download URL | tensordict_nightly-2026.9.24-cp312-cp312-manylinux_2_28_aarch64.whl |
|---|---|
| Size | 593.0 kB |
| Tags | CPython 3.12 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
eb16fd4607e3eb119ba91ad161c84873ced4e32019c928719a645a571b592ed3
|
|
BLAKE2b-256 checksum How to use checksums |
0074b33c1196ad0b429537fa5f555ad40814296f0b54d4bcf440edd77bf83b06
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.10.12
|
Release files / tensordict_nightly-2026.9.24-cp311-cp311-win_amd64.whl
| Download URL | tensordict_nightly-2026.9.24-cp311-cp311-win_amd64.whl |
|---|---|
| Size | 653.0 kB |
| Tags | CPython 3.11 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
4b044cf6ccba355ef01560c0dac92a9fc94014fdec8f9f397d183d200f151942
|
|
BLAKE2b-256 checksum How to use checksums |
85539c825d4e19d11bc653aba29a7fea9621c56de281c1576492380dbacb1231
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.10
|
Release files / tensordict_nightly-2026.9.24-cp311-cp311-manylinux_2_28_aarch64.whl
| Download URL | tensordict_nightly-2026.9.24-cp311-cp311-manylinux_2_28_aarch64.whl |
|---|---|
| Size | 593.4 kB |
| Tags | CPython 3.11 Linux glibc 2.28+ ARM64 |
|
SHA-256 checksum How to use checksums |
09b9a244e68cf36e9dcf68827b388f7dfdc37c64e98c79512446c106cce9f4c0
|
|
BLAKE2b-256 checksum How to use checksums |
d5835c61ce04ec3e211adbdde3e5c02a24e59db897236d01ac7ccdc306495ef3
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.10.12
|
Release files / tensordict_nightly-2026.9.24-cp310-cp310-win_amd64.whl
| Download URL | tensordict_nightly-2026.9.24-cp310-cp310-win_amd64.whl |
|---|---|
| Size | 650.6 kB |
| Tags | CPython 3.10 Windows x86-64 |
|
SHA-256 checksum How to use checksums |
32b014f657e223f9d60a190ec0420354bec87c0a9de1dea04bda2a81324a7da3
|
|
BLAKE2b-256 checksum How to use checksums |
57b60b59353644c58786a53a6165b8dccb13c2d8686fe67dad40226dfb2b3ba3
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.12.10
|
