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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.16-cp314-cp314-win_amd64.whl (651.6 kB view details)

Uploaded CPython 3.14Windows x86-64

tensordict_nightly-2026.8.16-cp314-cp314-manylinux_2_28_aarch64.whl (589.6 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.16-cp314-cp314-macosx_11_0_universal2.whl (578.8 kB view details)

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

tensordict_nightly-2026.8.16-cp313-cp313-win_amd64.whl (649.7 kB view details)

Uploaded CPython 3.13Windows x86-64

tensordict_nightly-2026.8.16-cp313-cp313-manylinux_2_28_aarch64.whl (588.8 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.16-cp313-cp313-macosx_11_0_universal2.whl (578.6 kB view details)

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

tensordict_nightly-2026.8.16-cp312-cp312-win_amd64.whl (649.6 kB view details)

Uploaded CPython 3.12Windows x86-64

tensordict_nightly-2026.8.16-cp312-cp312-manylinux_2_28_aarch64.whl (588.6 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.16-cp312-cp312-macosx_11_0_universal2.whl (578.6 kB view details)

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

tensordict_nightly-2026.8.16-cp311-cp311-win_amd64.whl (648.5 kB view details)

Uploaded CPython 3.11Windows x86-64

tensordict_nightly-2026.8.16-cp311-cp311-manylinux_2_28_aarch64.whl (589.0 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.16-cp311-cp311-macosx_11_0_universal2.whl (577.8 kB view details)

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

tensordict_nightly-2026.8.16-cp310-cp310-win_amd64.whl (646.2 kB view details)

Uploaded CPython 3.10Windows x86-64

tensordict_nightly-2026.8.16-cp310-cp310-manylinux_2_28_aarch64.whl (587.5 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.16-cp310-cp310-macosx_11_0_universal2.whl (576.0 kB view details)

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

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 01e49771a81fdb955d0068035088433e6ce80a078c0df72b4ef595f5e91602e6
MD5 742bfbd7a3f46805d9d0e719b073c639
BLAKE2b-256 571ea55a35e4823118ddc51ca3c56efbb75aeaea1c02d05913494b77e4101bbb

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 082ab80dbb54404f2dcce23a65eab7944df345a53bee484565278f477ae2b7cd
MD5 630923700320b65d086767df51b3b451
BLAKE2b-256 93e65c84de72d498703c937661a5621bc95e0f536d18d8a5424f6ca7a72e60a7

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 819a39ce5441103d491ab195b20a82b890713f8a60456e9eb1f0925b6a3eab73
MD5 d9def7f05c8dd133ef304614f9761f1d
BLAKE2b-256 e1b00e7abf4486f70a2d481ac3885885f8b55b33f5d643f9fe049af4fef45bbf

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 55c3cc0c4b304cbe0858e978b068bb55e5580fc9b38dcab14cf2ae635c9810e3
MD5 77f40e99ac2ffc2b9028291b164b9fc6
BLAKE2b-256 aa408d0e9f765eabef0b2a7b91f388ee518ae28b6d3ab058d845d620925c2ac3

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 602416d347971b0b465b28131ea1274974b92a2dd619eac734005a37b689a9e2
MD5 ecade9adf4f816649b6b01afb8aa0a4f
BLAKE2b-256 a5706e6242ca3e0cff40d5e658f842f2034faf1d7842b4af8468205b27e96533

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 62d617d7031b2a03fa5340b3816ffc7720bc7777812f09e6d6b13ff7d870dc37
MD5 6b5c3f8e7a92c10c21d6a62a01d1efd4
BLAKE2b-256 055eae941fd71e3ded915b8f92cdaddd97b32536143479b5e6233e78be2f6da2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 331e0dcec00973e6ebaa1ed34214bc79ba70b72a98d999f23807c2f5dfa238da
MD5 a0b0e8bf50a781bb90e4c49cb9e71d15
BLAKE2b-256 89d364847256954a46a463325c8ee63d9d2834217c67f8ef586a3b1fb493c35c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 51c756c125b6ff14e277565b618e38f96f805e27751901a31fbe1a8d4e90ef66
MD5 fbd9d356fab357543c43efb9e4b31607
BLAKE2b-256 7cc5d0911d90b56bd9f0c5e4c9384cc13842198be70ad2ac90ad75523cd46e47

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 ccf2f78e870f20cff1933825319b7ffc3723bf8492653b6dbc69a18a4f3469f5
MD5 bfcea228f7f096bf53a8829f5b1cc6db
BLAKE2b-256 2efb18430506410aefafe45b5f5684ebf624b7cbd4b19679fc5f207d833c4101

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 6498f50bf864c1d10e505f553758ac14c95c13abddcb90b385731ed62a4ae660
MD5 5c8386519dd4ffc56fcb1d557e179ccb
BLAKE2b-256 dc8d62bbe43c06dde91457b13ba8c34213b0f14a85d5e84b65a81f8f42d6706e

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 2edcd94c1c3ca773050444e56176c7a3b9d29c2d485646c2069269010263bd10
MD5 478a22cc7d9be56114757ff215adbf3b
BLAKE2b-256 f325d7a6b34c344defe790da8282e5d7d4e6f22d016813f1b33b1bfe173852d9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 116fa70a9b76417001c7dd7b2a5605f0224060df37948e82ff075857c6af1f95
MD5 39afaca2c5df19544585f28e45c73b8e
BLAKE2b-256 38a7c009fb31e50e86a7b2b864d10f80cd04e9fecd890f930786a14d4f1be819

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 60d9bbf8c7440d33609a2f248bc09596cbd327188131befd6a95f3cb629ea1bf
MD5 f7502db6a74709a0493f5dbd6aa2dfcc
BLAKE2b-256 caf029d5810ce6fd27ec58941e844fcd0c45f23c3e78b70f1f82ee5fff6d06db

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 bd7d000070d88319340dfd4223ef799f800bb268f05543f57f1fc0e3795f2a2b
MD5 4a17921ae2f25254b96389152b929b2e
BLAKE2b-256 1995b21638f8e4a592a2f2f82579abf9d6da8fdf91e6540732716f2e2eb6721a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 1b9dc85758f7f3eeb60300f41bb49e6c5a9405d4b7964974205ec4dfef2e8dc0
MD5 18845c55e8f24c3ba70c24a99e418f41
BLAKE2b-256 5dd0e5ec2aea0b4484884ac270906a0d43565c8deae2e00c683089a73d9a3b5c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 5028ee685c8e7a1bcb27aeb3ce131a9c7da9345bb526ed380405060de0258e7a
MD5 cf15407eed2931f5925813f5740df0e5
BLAKE2b-256 78abbe338a62d1de488a77bc1e04245f35d09f2887de5a918498b04f3b719e8f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 2049bba7df7163fb41cba7429d0eff5f979897e951e12fd6c6597f9368fd17c9
MD5 87dd2079f60cb0840a62daec4cf12a47
BLAKE2b-256 251e5b6cc51d78bb366310b3f70ba29aee74a901dce96984b0045a8b2bf22c21

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 7611ddd7466bf3cb3248f6ced7aff1503921aaed11988ce564d434d1f4320548
MD5 9052bd64f940010c752aab475fe3f3db
BLAKE2b-256 163c4a52866f02938c754795711ac402941f8ca0654816d8622c3681ae23993b

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 6d8ff2c361f3da2ec8c72fcea4e5256e0def28d0a354220d372663315f45d225
MD5 f454cf36ac881890cb99783049b2f30a
BLAKE2b-256 266e27e8975137e186822beb1d45afb4ee6b1880d5fc4d46cf30c40c00bed3d9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 4e3e99937b90e8636854c81b31c49b2d487735c5639adae680540d460840454f
MD5 c17e9efe4b7b79c28965e2945e7d126b
BLAKE2b-256 4df913714f566f3221f7aff8090c55cd8b60b7cbc2d2ae449a9220c875268311

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.16-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 287f48d8d3103fef4a48468f4bce617a67bd64820117d7ddb22d9481fe772e3a
MD5 e7fd4f4d7bca0eca7f900921505318eb
BLAKE2b-256 7341b7aa0ee8558bdfe37c67967cb5b34ef6248f608a9c088c07ab7da0694c47

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