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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

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

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.13Windows x86-64

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

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.12Windows x86-64

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

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.11Windows x86-64

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

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.10Windows x86-64

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

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.13-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.13-cp314-cp314t-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 d1d2ecf20a6b2f900ec2be2abbcf4d3d0c5ec58fdc1f449616aa7f0a7ab0b229
MD5 118403d0fcfbd014fa22a1cce2b8313a
BLAKE2b-256 6070d7cdcc1f5dc61631cc1c3da605fc9b4f5ccc2b9236a8df47ed5df5caaeb3

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 dd01471020ff82278de46acbb9e141934ca112b64a5b5fa6a5d4891f397984c8
MD5 80f8fdd8afb6b3f3bf04cff1e4648cf8
BLAKE2b-256 d90dd99712ba53d8f7b5d2a54dd20c45c71aae845000022cb26c59940bd81aed

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 a2fc4311701f8109762872e09ba26d3d996e3844d7318517d342382fb5627e5e
MD5 79c428822be58e3c075979215db409fe
BLAKE2b-256 0492e259b2efe51b8f1459d8e8f0c668660c3ac8ed040c85de30f1b9fd13ad77

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 75f4aa6bc696b5a2d5d67a5df3f595d2776eae9ee6d17719e02143c56484f2e8
MD5 d52366aa32d2ea321211fe069a09e4a1
BLAKE2b-256 eeea63ac84fca1713da1a6744338b53ee7b3316f47529d28ee25786956ddb2c5

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 f6433f242e7c91fcdb9fe52093e1bdb605c08b5754eee5c70b8cb4dbd3cff09c
MD5 c423c24f7ae4915e0ff73c5e10362196
BLAKE2b-256 8135b030187d4ee9bf9b3ce74d1395f6fce228143c82c576c46ce6a33c1726ae

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 fa73fc1e00e9b534cade123161ac35983f5457b8a58852dce334a2d2439f7895
MD5 e7c6cdbaeec52176494192c382cfaf04
BLAKE2b-256 071dcf6cfc4b05cbf7ee44e88856de13d5b0e558f730b2dc9e783c91aaee4d5e

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 fd061c4c9b366f9c84fd8552689332f2401639dd30744e79c7f3f5e2ace91ab1
MD5 4cbf9e4df57dc33f0dede8568e876947
BLAKE2b-256 0125d795ade18c9a9412636011522337c8706359fcaebbb1dbb7d80e3ca1a4a9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 d96db47afb0a534c10185de42e8f5755e1cf5837713f238f58e0704929c04d8f
MD5 57fc3b9d922550b7fb1810eea0e9bd7b
BLAKE2b-256 6147f47a89e5c7c01b9bca0ecad1ee0da4385c8a41a703dc483af202e7d44d24

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 64da85ccb6156bf8c6d99fd283fa74dd75b3601e0494651b9edb6b33a136e788
MD5 ac2090f7f2d50c1a469dd8b9843a64d3
BLAKE2b-256 39340f9630046dbbad959f6c83a3649c25cb78cd3b7494a3c3d8a87b21dd3728

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 5205dd49503d5fcc01be5ec069b0d212e514f0affe81c9e3515b9b9ec4e1634e
MD5 52d807371d066f50b6dce0cd524baef2
BLAKE2b-256 40971d651d7da1aae9e76e6f2620d6acb010e0e4cab0a514689ca54fe9e61922

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 bb9b998c1a4cb940eef8218697a91fe248ffcb34f1d6956ca0886e60ab0169d3
MD5 f8d9408e4f44e2c3297edec8be294b38
BLAKE2b-256 813dfe83b3a93a7479fbec8f370d0a4ed3ee859e09008123d87860008abb6d05

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 7f6eb12238e252b2ecf2680827b8295efdc5076d076e13f38b57f1e6aeecbdea
MD5 ca344b611e56c12c50ed0865d10f34cc
BLAKE2b-256 3d5189eb0ae8d2c87f6f287527e9a48eea4acd42efa386dbf475a6db526bc911

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 6032a40aaf6908a90afa2707029ce240e63571a846c22ba95a4b8b26ea252b44
MD5 b8da95d4299879959a4821421de1ab0b
BLAKE2b-256 009381e0284a95c5a4ce280cbfd2d27e540c72d8aa7bc01816d42aa009930316

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 a67b3116d4ce5afc3158224ac33a5ce85b385b7966baab174d4eb73fda42a849
MD5 ddcb2af0ea170d6cd00b1c34b1a691f6
BLAKE2b-256 d0c9c07c9f9cb2934b8cf1e374edf39df7641066d8685ea8b696cc41c23a890c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 f08453b80b3bca39d5ee2fa8a35411c9717712c8c6582b32d94ef6ed4f562eaa
MD5 de92f5b1e72c527c759b37f9dc5be5df
BLAKE2b-256 7e3a7edf02ed72bb1d54f9c801a7b49d8aef2517d009090bb4685c5d2d5279b9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 fc93723d22df17748d01642d7a38621e84382a3e06f0a87fc8f6ee0b1d5159ee
MD5 109e6dd3c9975439b93236e1ecfcfbec
BLAKE2b-256 deac01793dc5bfed1164d2281c433bec75cde132adc868669c58ec1ce9827f20

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 029371b3a1a570bd600c10bf5eed6fb8986f64ff6b196c4ce314164fddbbefca
MD5 02f55c340088b591a82a8b6c73e93a5b
BLAKE2b-256 b14d3f37c4d2f5a8ed18b6770d4ba9fb4bea903e4e5b1a3cd87f9c66c6f22ec4

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 2a0434ec16cc9c20a7f315580d6f5117c679fb5dff088b3e191161ed5da920cc
MD5 321a836374b3d580539cdbc8c102e412
BLAKE2b-256 7fda90fbec23bd876fe836f6cf09cbf6b032c4404734fa919f8e6477ff57020d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 cdd4db0bb09ef6f66a4048352b0a4f990e592b1d9cf86a10b842daaa5043d2b3
MD5 dd92bf549e4012f2c0a7b63252c92fa3
BLAKE2b-256 2dbe983ca173e732ec4c694d4b0e035adb5009821e7979450be73ca0e1fa89f4

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 fb58408710f84b536eba6eed7fd05e3d804d3700616921cf9347e55fe727dfa1
MD5 3aaf00d99a5332488f5a2d662b7410f8
BLAKE2b-256 a6427e79efd1c25354e8ef34420bfc378731436aa444c6cef5f02ca24d594453

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.13-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 a3923c8ebd28e36f55d6184e6614c8342dd014d5e0b0b9d81483283a3fe611e4
MD5 d72fd6f9998be87d3d372f5491b86685
BLAKE2b-256 d23a76dc6a5ab7c4d97e838f51341e5eb421127d1aa47c33f8ac9ce4f2f414f6

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