Skip to main content

Docs Discord Python version GitHub license pypi version Downloads Conda (channel only)

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); pass non_blocking=False only when you need an explicitly synchronous transfer.
  • Shape operations without boilerplate: indexing, view, reshape, permute, unsqueeze, squeeze, flatten, unflatten, stack and cat operate on the batch structure rather than on hand-maintained lists of leaves.
  • Low-allocation workflows: lazy stacks, preallocation, memory mapping and inplace=True shape-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.compile coverage 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=True shape operations, including gather, repeat, repeat_interleave, roll, reshape, flatten, unflatten and contiguous.
  • Improved torch.compile behavior 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_module and compatibility with torch.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.

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distributions

No source distribution files available for this release.See tutorial on generating distribution archives.

Built Distributions

If you're not sure about the file name format, learn more about wheel file names.

tensordict_nightly-2026.8.9-cp314-cp314t-manylinux_2_28_aarch64.whl (589.0 kB view details)

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.9-cp314-cp314-win_amd64.whl (649.8 kB view details)

Uploaded CPython 3.14Windows x86-64

tensordict_nightly-2026.8.9-cp314-cp314-manylinux_2_28_aarch64.whl (587.8 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.9-cp314-cp314-macosx_11_0_universal2.whl (576.9 kB view details)

Uploaded CPython 3.14macOS 11.0+ universal2 (ARM64, x86-64)

tensordict_nightly-2026.8.9-cp313-cp313-win_amd64.whl (647.8 kB view details)

Uploaded CPython 3.13Windows x86-64

tensordict_nightly-2026.8.9-cp313-cp313-manylinux_2_28_aarch64.whl (587.0 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.9-cp313-cp313-macosx_11_0_universal2.whl (576.8 kB view details)

Uploaded CPython 3.13macOS 11.0+ universal2 (ARM64, x86-64)

tensordict_nightly-2026.8.9-cp312-cp312-win_amd64.whl (647.7 kB view details)

Uploaded CPython 3.12Windows x86-64

tensordict_nightly-2026.8.9-cp312-cp312-manylinux_2_28_aarch64.whl (586.8 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.9-cp312-cp312-macosx_11_0_universal2.whl (576.7 kB view details)

Uploaded CPython 3.12macOS 11.0+ universal2 (ARM64, x86-64)

tensordict_nightly-2026.8.9-cp311-cp311-win_amd64.whl (646.7 kB view details)

Uploaded CPython 3.11Windows x86-64

tensordict_nightly-2026.8.9-cp311-cp311-manylinux_2_28_aarch64.whl (587.2 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.9-cp311-cp311-macosx_11_0_universal2.whl (576.0 kB view details)

Uploaded CPython 3.11macOS 11.0+ universal2 (ARM64, x86-64)

tensordict_nightly-2026.8.9-cp310-cp310-win_amd64.whl (644.3 kB view details)

Uploaded CPython 3.10Windows x86-64

tensordict_nightly-2026.8.9-cp310-cp310-manylinux_2_28_aarch64.whl (585.7 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.9-cp310-cp310-macosx_11_0_universal2.whl (574.1 kB view details)

Uploaded CPython 3.10macOS 11.0+ universal2 (ARM64, x86-64)

File details

Details for the file tensordict_nightly-2026.8.9-cp314-cp314t-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 db26980cbf4f1efabb2bcf145fff2f8579e49e47e433312ca662a4da779c5e2c
MD5 07c1ed8ec83aad5511bcebe4dc5cca7c
BLAKE2b-256 ae5352899a55757fb9968b584e3ab7462a43f4fc63288b4bbe4a47581e6dda33

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp314-cp314-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 dc6a53cbefb1e3fc4dce0f992a400b2f68f41a7cb601c52cd951e9c840328c79
MD5 5bbdc9a20409f0cf2c29a55424d7f37c
BLAKE2b-256 a01fa9c0b672f34c871800845ae642ad1e8ba8dc03a102c98344d7d1f4f43e92

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp314-cp314-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 e8f072df15092ab136a3c6f097a8de0adb113af7357e834d275144e140e1560c
MD5 5aabdca1cdb9b0c53070a3169e184aac
BLAKE2b-256 188030d4323b1a0a51b6ac1ae5332d13a21fda02fc380368d8d5948eaa164066

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp314-cp314-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 cd1230dc65a7546291d82c4122c647feb11aec01cd8ea4a2220e14f07d00bd8f
MD5 69e37d2a904d14cfd138cee1cc17a111
BLAKE2b-256 2e3af8934ffb08df98ce8ff540dbfcb68456b579c6aa2ec7b44007335ec4f78f

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp314-cp314-macosx_11_0_universal2.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 54f612751cf90361d60b0a3d32f4321bd7033ce710fab9cc9763a6313a37fecd
MD5 1f26e00af1d7fcf7e177c21f7ed9fb00
BLAKE2b-256 c87888c30a778a15274bd6039fe822f1ea059b76b916a48fb445a2e4a60e029b

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp313-cp313-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 278b99a969cf8cbf3ff877fe6f6ba674ec85e524c67dbd54f416c6a6c6921083
MD5 35871997146e4ed5e3bd2c50cf66f03a
BLAKE2b-256 11998e5b43fb3981587b567eb2004f3a8f5de666c83a225b850073dd6027148d

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp313-cp313-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 1fefe91484949c56cfb528fb84e8b1b13df81e299dc008c279ca908c93e169e1
MD5 606f01f5c93146083515e321f6d4b7d0
BLAKE2b-256 26f518ceb3b4841f156f6a6e9c6011a36ce542e79d0dcb1a9171aee54c300ef2

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp313-cp313-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 859d1493c6c229ced7e8750783903059ae6b6df2f4c2a4e1b1c9d0b44c76e941
MD5 8aaa6296e4fb11922c5e2b6c3460cefb
BLAKE2b-256 8dfd566829007cd15bd93e554aed54d80fd2a07f5216793fb22d347e61e13e5b

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp313-cp313-macosx_11_0_universal2.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 ddad6f51cfc91eeaf9098b2e0b981c1570741a7fbfb943fd4cc0efd448b77c62
MD5 4a848cdf1ae84430967ea4d00e8b41c9
BLAKE2b-256 12787e9e20b0ca831f7df5f27a27da845425f8a730df5876edfdf0893b3b8548

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp312-cp312-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 eb63e11a5ddc61324bee470ccd67233a9ef5bcb252c4734905b387d7babc3e81
MD5 ee67886e2d96349c427f4239509d3e80
BLAKE2b-256 702f6cefc247bb8cac7740d6178b11164b266da2bb3dfba856e4773d822fad64

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp312-cp312-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 f2ebd1d78242baba544dc22dc9d79d14d9b5678875589fb732a15c0de06e11a3
MD5 1e23a751966161f5a6efe2e19d5700a7
BLAKE2b-256 008fbdbda0faaaef309e3c59df2324ec1828d1848b15e476616f7b14013e6a31

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp312-cp312-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 cbed1651ef8736b70c4c7cc0c63b0be32a92c02240f9e5da80a1c55a42fddc5e
MD5 02914b667bb7d51d594ffeadf9fbfb8b
BLAKE2b-256 36eecf12a0c0bf30346536c4d9efc5b73947740264a5eec6a1065ef93b06cfa8

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp312-cp312-macosx_11_0_universal2.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 1271fea93ef766dcdfc357887fde090a8b6938fa005344da7511898183f53ac6
MD5 e0a2827ad6af38d3b949401f1985c4ee
BLAKE2b-256 17875998840b570ba729bceab57f54023e321d5a25414383b3e8d2b791caf9a0

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp311-cp311-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 ae446bedb3ace81dcb4021bf923b2b869d410e5705f1366fc9f369ca7c070370
MD5 f7825f0693ae2b31285957ea23b29027
BLAKE2b-256 fcf6839463bf97948c5d026c164b0dbe5c31ca717cfb71f239a64f6f4ca240b0

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp311-cp311-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 c02a60d0b6a9bde58eaf2d8fb4ddc48196c822d366a75cac9d260be87de43323
MD5 8a92d7f97d3da6baee151c1433b66651
BLAKE2b-256 11c944019c35ba90e9dc900323f65fd0e3a3070f1686ccd616278ba448a3ab0b

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp311-cp311-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 9a52a379323aafde264f9e8d02210b9b025794c0c33149f6c5cd96c206d10941
MD5 b512c5f3501d4b8e213412afd636ad67
BLAKE2b-256 349db69318b4bc0bf31217ebcfd4751d36558590a0dd48f8cc286f66444b909e

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp311-cp311-macosx_11_0_universal2.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 c86a980e8bdcef38441a1fc843227abbb4dfd2d5e6d93fff8bceb00aa53784de
MD5 972edeb68ae47be3919ab004adce5f7d
BLAKE2b-256 c423c8936954ba33531970aa63f63a47c46973459722e27c166e47d094fae8bb

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp310-cp310-win_amd64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 614e9bc70384c0f1643dbd6b1069e12674e10787afdb4d3b4ff1ac0706b568a5
MD5 f9a67dcb363030ff1130518cd2d9ecf3
BLAKE2b-256 6f99c892619076b2a5b0ffb03434ea0cdceca1aa006ada98b5c662b916ffbf1d

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp310-cp310-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 e8b6bbcfeee563ef02bafe902ba7b0fccf0c7dac79afa2e17e0ce51f134e25dc
MD5 e3bb79e6313aa7eee83c2edc8c42a1e2
BLAKE2b-256 419184a898cfc168d1efdb88ddb319a334875c8b4aac4a1d182c11d86f1661ca

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp310-cp310-manylinux1_x86_64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 f425f70e1e11120a370da0d4d56ccbed4bbced83e2259c291a38fb94cc10743d
MD5 408a98d953ac265f94ad759bec733223
BLAKE2b-256 89ad98bd1ff5f7634d2c11ccfda4631fd9c9a68bfc0724855ae54dea5f83fcab

See more details on using hashes here.

File details

Details for the file tensordict_nightly-2026.8.9-cp310-cp310-macosx_11_0_universal2.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.9-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 a8ae54db6dbdbef657d785304695ccd24cd86a348e3345e67ec2798232b166bc
MD5 64dea88e966332313cded5a9a9b38bb0
BLAKE2b-256 e76c813e2028f8f35f7ec63271008d05355953ee40c352a722ce1ebcadf12419

See more details on using hashes here.

Release history Release notifications | RSS feed

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page