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.11-cp314-cp314t-manylinux_2_28_aarch64.whl (589.0 kB view details)

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

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

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.11-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.11-cp313-cp313-win_amd64.whl (647.8 kB view details)

Uploaded CPython 3.13Windows x86-64

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

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.11-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.11-cp312-cp312-win_amd64.whl (647.8 kB view details)

Uploaded CPython 3.12Windows x86-64

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

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.11-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.11-cp311-cp311-win_amd64.whl (646.7 kB view details)

Uploaded CPython 3.11Windows x86-64

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

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.11-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.11-cp310-cp310-win_amd64.whl (644.3 kB view details)

Uploaded CPython 3.10Windows x86-64

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

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.11-cp310-cp310-macosx_11_0_universal2.whl (574.2 kB view details)

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

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 3e593db4b552735ffad4d4343d32a98080b59dac3a8b4c7f95b6a20eda1558eb
MD5 650008951f31158dee0c3055f5fbbccf
BLAKE2b-256 e67da9c703e298f58519d0aac702173fefc3f308874b9cfd62e7d8d2716650cf

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 db12decfa04fe89777b2800d3fdd9a31fdb143f434968859635b61f8b25967ad
MD5 6c3900f93a6ec13c62084ace330912e1
BLAKE2b-256 47931bcc875723dadd2ce1f14e19df83d67f20853e61eaec04325b69e04e45c7

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 3ad583354299aa54cb5ed12cb5b08f21da3753bffa2ada27cf31c59e56924905
MD5 1db38201f79091a2da9d2d00f81cce53
BLAKE2b-256 b729308e4146bd6c2ee3009a0800075c9f370a6ce56d57c421a936167b0a5a4b

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 45cc24a0ecb5c1d75bf18396d3282bcfb55529217a99d108016dbd97eefc8bf3
MD5 8cd876d9e2f4ad25a17abc6cd7421889
BLAKE2b-256 e57eac62975fe1a8ed804cafd7eb36a9c1302caa5c3051e139fcab8b93649b39

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 72849bbee4f05c6cbbbcc2863e00b66b554b0377c91227bc5d4e5753a8fe37af
MD5 a47c94c2d73b44f8acdd9e8fd63d1f75
BLAKE2b-256 1dea04d2924edd77badd46945fe00cd77b486a0983030e1ffadd362077721261

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 553742f152673f713cabfd789858117d96803019a7daca88c99dd7566da90e3e
MD5 fa2d5ce4e08260e77687efa6646b4e4d
BLAKE2b-256 43cc6e5a740bb987a8cd6c4d73564d47e79cdae003b5017c41a990f93164c6a9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 a49984549c76af506b69729c3de5de35e80d754f48611910f5c699520b89f5cf
MD5 bfcafa7cd28ab96cb4caf8b8d8de3849
BLAKE2b-256 f322aa6ed17455ef52f01a5bcbcf66c3e6614905670b4607c9ca6f5ccaa98493

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 95685ac0aa872cf1044bda0adbb8eca897541ea7af4d70499ab2aec22a87d87d
MD5 997e1aba091729cafea45b44b637a5b6
BLAKE2b-256 fee6fac78ba7c2b2d3c707057028b219ba639d205003ea8ab47b60e714b59ad9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 91495018dd3be0eaaa0625b35fe3f82abee46bfb6c14506ec84ceb73db68a570
MD5 657cf5e6fa2b20a8dc28be24624a7a4d
BLAKE2b-256 eca0a3139f2bc389211c5669ba0a8b5039edf926499676fc73639441a430d442

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 81968f175409e2f0fa4a028f36437db6cd845c3bbf41178ebf307d92fb8c29a2
MD5 bb33a01493c83897851bb48cfdc987ae
BLAKE2b-256 56f36e5f96582e5a4619afc5f3aa3d49300252f6501ea14c8f21ed8225705c08

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 2df166df7b32c83866c1089e3b22307078e7f7dab6635a4fa38a0e226968238b
MD5 6fa4761740a27b7a44bf13fdb0f67f3e
BLAKE2b-256 41a256dc9805f06c8a7f2bbe6c57ac212d43f199b96d437509748a9cf178567e

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 b3afccb5d68caf100a080658bede77dcc6a969f3a4578b65909a03bc200ae885
MD5 641e7e31e8af43caece3cd757dd218b5
BLAKE2b-256 ac5251ef8d7b09d233fd28a9c30147d0831b2dd532d2fec456cd36bf3c78245a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 ce1121c3b184e335ae6cf9621c272d8604127c004236d027627503d46a6f84db
MD5 0375001e4582297ba6fd268047cb53b7
BLAKE2b-256 7c4ed6e5fc4e473f016e8d97575d28c397a8b2abce03b69dc02ae17497c72a76

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 ffaf9833aa52c3272643cd445604db8a75e070ef458173d3de02a69ba0c36fe6
MD5 1c340700bf9e20f0b8443cd18b7a179c
BLAKE2b-256 bb1180b04a1ae228a13e2f7f22ed9bfe844763d5b31c95252b2723349f511653

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 961047224446cf34dcd09db3024ade302672bf13bb83d3f13f9faa15024ea276
MD5 28c3723a6fe7a5907673b6288d317bd8
BLAKE2b-256 6420e286ffb59a5252e4f163cef42a94fcf5318c521356a7fdafce003f42bc59

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 9f2e10804355cf110503267528a7f3842ba5d7af1fa9be092bfbdfb9934886c9
MD5 833a56a48e27ea9fdfcb42f31dc8e2db
BLAKE2b-256 8050f812c6f7ca0072e840083a8080a63f245b4e0171f8276607750b25645387

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 dc7792b370af6c7e48a98ba2e0ea7bd77d42990403f44e392ef3859d93bff8eb
MD5 0b95456db7db892a1a391e0d0d9ffd9e
BLAKE2b-256 45caa3f491985903bcf7d7bd3160ef5728ab19baeb337ac12a531bb2a5244df5

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 9bce112264ca8d8ec28234b046f2d47109897a0642a4ae3b589fa6a7d079d3f5
MD5 3a8e9763b28f8c7dd05310fa36a9e4a3
BLAKE2b-256 ae0b5a992318122efdba7538efe2974a7d7b77eb11a5fe0d4c2789fbdd0db94a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 d109c31cd617b826384c863431b263814c67166c195cf73e77b960bcb0633134
MD5 a610cf620c9757fd86f55bf4e233a64b
BLAKE2b-256 79151a844e4b41af3789053e7235609dc02f28f043fce46c396aff5ca1c1b23c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 a9302e8a5f8147a81d1bfabbd19e380e6e6efa46f76c2e3f94c5a143057945aa
MD5 f309c97504f33e5706bbf9ae7a9015da
BLAKE2b-256 345a46f73420c45a70e6985d87bcd12843be6858759ee8b0c19ec8a849c28fac

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.11-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 5ea3ea54fa0b3c1470878d1df1ebde0e3eda91b8fac79f8fdbb391fb85cde716
MD5 2579fc051ad2c337ebb91c7035b841f4
BLAKE2b-256 df22a8ab5fcc101ac81588f3c7c08b8fa8e05b05e7e49111987f1bc4f741d4f7

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