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("/tmp/batch")          # memory-map every leaf
reloaded = TensorDict.load_memmap("/tmp/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.8-cp314-cp314t-manylinux_2_28_aarch64.whl (587.7 kB view details)

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.8-cp314-cp314-win_amd64.whl (648.5 kB view details)

Uploaded CPython 3.14Windows x86-64

tensordict_nightly-2026.8.8-cp314-cp314-manylinux_2_28_aarch64.whl (586.5 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.8-cp314-cp314-macosx_11_0_universal2.whl (575.6 kB view details)

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

tensordict_nightly-2026.8.8-cp313-cp313-win_amd64.whl (646.5 kB view details)

Uploaded CPython 3.13Windows x86-64

tensordict_nightly-2026.8.8-cp313-cp313-manylinux_2_28_aarch64.whl (585.7 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.8-cp313-cp313-macosx_11_0_universal2.whl (575.5 kB view details)

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

tensordict_nightly-2026.8.8-cp312-cp312-win_amd64.whl (646.4 kB view details)

Uploaded CPython 3.12Windows x86-64

tensordict_nightly-2026.8.8-cp312-cp312-manylinux_2_28_aarch64.whl (585.5 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.8-cp312-cp312-macosx_11_0_universal2.whl (575.4 kB view details)

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

tensordict_nightly-2026.8.8-cp311-cp311-win_amd64.whl (645.4 kB view details)

Uploaded CPython 3.11Windows x86-64

tensordict_nightly-2026.8.8-cp311-cp311-manylinux_2_28_aarch64.whl (585.9 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.8-cp311-cp311-macosx_11_0_universal2.whl (574.7 kB view details)

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

tensordict_nightly-2026.8.8-cp310-cp310-win_amd64.whl (643.0 kB view details)

Uploaded CPython 3.10Windows x86-64

tensordict_nightly-2026.8.8-cp310-cp310-manylinux_2_28_aarch64.whl (584.4 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.8-cp310-cp310-macosx_11_0_universal2.whl (572.9 kB view details)

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

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 9841863bd37d8837549cb225dcf3aeaf1d4e0c590099ef0ab36b635932facc3d
MD5 910f04db69c9179ee6b2afc63e454f33
BLAKE2b-256 f0c68ac0f7274a579d6f65aa7c9687a00d555ce744851a25b603c87407a2e10f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 c475ad89bd65eb26d68658fcfb2225abb74de5bf4c8f7e7608c3b4b9b7105a56
MD5 7112a51c644886c3359d1ca4d1dc1979
BLAKE2b-256 d7090d30f7f2d1e05ae0123fff5f6e23fdfb03195523bea0363ba03558cae70a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 f575097e703d68e4de91c8f2d07491f319fc3d966e5dacd10b9be1d2248a310f
MD5 bbaa51ab3b1ca28645fe316c502cb07c
BLAKE2b-256 34390f78a14b956d2b933d18d92a5e431bd727a5be3cd74e854daa08d032c969

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 b2cdb99dae0f4157ccf0a321936ed40bfa349e1d7097162ffb3e20c6668698b3
MD5 3d166824ecaf277fb2e7138b6b0cfd2b
BLAKE2b-256 1f81b30c591a6d499c57ad4f656cf572b0b6a9f9a911a1719d384bb51b2513e0

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 59b27e01be6dc1b89052cb47ed540f0a2876e299b4c8173862136309b0ff38c3
MD5 9e0e2c3c81431cb8ecd852c8d4bb1ded
BLAKE2b-256 15ef46eb1e32aaae39bdd460f6b0bead986f8df11bce49cca43bb1c32411af60

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 2de16be49fe087ca31dc9a89fe05ea382275b6f62513490832ef6a334e8c14f5
MD5 db29599e19bfafc002a80be00f0fe002
BLAKE2b-256 9cd881b2859e1731e8f92e571213558fd43cf3b4c725d6e641032dbe10a77a9e

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 bf3b6509cb5ceb3aba70fa7b8317dfce933673152fef52ae177451901515b386
MD5 9ca88b84160859bf78c9a44b08334654
BLAKE2b-256 d75f0c2d5f243e65d1c56c27228026e72beb1b4c15bd082c6fca5c90154e79c1

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 2e08652c876f51624cd732ba5c66fc3292c53280400e00c04aa0f32a9484463d
MD5 b9bc789a8d5a17366e26b520a41fc002
BLAKE2b-256 a6b7660905c20850b3123096dd6c6064c7d4ad0b7a148ca640f82dbf28fce51f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 d82566849dc2db03284d8367ed32cafeeb756a507f3823c2711f714a429261fc
MD5 198fa92a1edd6552f33931c47b9262cb
BLAKE2b-256 4c6fa856596182ba4efa2c555c7874c10184b9b6af9cd1483f93b7eb7e3e0d9b

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 7382bc55065479ecb9bd72b51a214b2cf9bf0a25cd4ecbd26a861c91cc7a7364
MD5 dd0e83f90c79db7bddd7872bbf49d179
BLAKE2b-256 08574e5b303062daab68b680ed4cb2e86f295ab8ed28ce224fa28029444f4b49

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 817a8a4b62910d088c62a330455b173a4ded26b68fde6b8ccbfc005301133ed2
MD5 e9dfb0752e6ff261224fac0431568ee2
BLAKE2b-256 38d9b392e996bab843407a184d7d9db0f843f5d7f9c76752035c040aa5d26fde

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 887aa6b37a38a247eb39717ce452a296e334b3a54c2ff524f916e3d9466e87b8
MD5 e691d6d20fe0f58a5acaeb1f082a6ee2
BLAKE2b-256 34b6d434111a3ce51d987c4bd382eb2971c444297975a5eb9d1d6da8bf3911cd

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 b3e61745c7ac358ecc251464192cfb540491b1eac31fcb1ef0ea799ca77664f1
MD5 78d4e21141a6c57e6ce6750e9cabcc8a
BLAKE2b-256 3844dbe491586381179f17d4ac38bb1915a6eb623cec5c5282351b4e27569742

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 5d92cad90a768aebc32df3f3d6912d3824d2238c4751726a162a33efd739869a
MD5 5c8b12c4fce396218f05f454a1e57b94
BLAKE2b-256 dd14933f0c59822f0afc2de7048f383d11ad6a88a7ee0250ecf0c3b4804a8460

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 4413c94bdc70e276eadf406a1d0ca4c43954cf09d6ae10622db5928786afcb48
MD5 cf9bb72cb2d9d3f05c83a299bbc854ff
BLAKE2b-256 ea3f8bf827658831de5d020373fb698d6c43c6c2277fdbca3b75106c77104977

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 57d2130c11145baf81fba2a36b66ba5def48b465df5ada1c35f5ed48238aad2e
MD5 6452745b3448c7e4275c2c8f005040ea
BLAKE2b-256 1b9fd40c5ff8972167958be9870496b8bf259f960e02222f2755a76b2e31a8d1

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 0a39ec5dadd1558518833544a07683730f0fd49f59cf698ead3d320e0f9c1432
MD5 3b6427035f572aa142d07b44f4a6e031
BLAKE2b-256 061449dbb015c36972ecc8f3fe5d949d481584df21f52542f1ef5d599e24e7a6

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 32e1855983300d33faf5da3742e79d8625ecc70ccac28bdff90c9e8212da3f2e
MD5 c151fae3aed521d240a23d5a2bdac42f
BLAKE2b-256 78add1aa84cdf325b6deebf1c5bcd42da33f23126c6f7199ae06616be6aa3cff

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 182b55e8b998dcc869e71d0784e23c84b0b708a6b923fda5d9e07f319f5a4533
MD5 484d2dcdd70e18a2ec2c33c699f1c13f
BLAKE2b-256 8b8dd56bce2383ae5e81dac21bcdd4176eb84bb406777bb1c4f9da9b193d5dcb

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 68d864dd31446f5e258a57cf225b89c8899cb06f3c9ec67c64d9d18f21eff373
MD5 3d48088782efa98ee39e203efcf647f3
BLAKE2b-256 6b82cf191fbfa6d88083055c4560dd8273dbe480a89090c7af603c5d0fd424a8

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.8-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 b9622c93a9ee4326f5699c26d390080af8d784e1e2a7255a3f16cce125c7e9ff
MD5 8bf8d8f4e24be6b0e3fccc6e41b92fb6
BLAKE2b-256 c947c7c8a8ec69c00a8dca10e77a16512a7e64c0c8033f890cd5cadeaed67525

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