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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

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

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.7-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.7-cp313-cp313-win_amd64.whl (646.5 kB view details)

Uploaded CPython 3.13Windows x86-64

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

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.7-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.7-cp312-cp312-win_amd64.whl (646.4 kB view details)

Uploaded CPython 3.12Windows x86-64

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

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.7-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.7-cp311-cp311-win_amd64.whl (645.4 kB view details)

Uploaded CPython 3.11Windows x86-64

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

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.7-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.7-cp310-cp310-win_amd64.whl (643.0 kB view details)

Uploaded CPython 3.10Windows x86-64

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

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 fd5e60f41bbe18a9305d617bc2d7f734c3979a3dd1de2049050fc8b9684822c3
MD5 22345ba8742073b24c329d62197cf23f
BLAKE2b-256 9d566002e1de1071de6ab6b5a99076e820ea8879afeec117501a52a8833c84df

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 86aa2f97ef66e5c0ecc58e46ef09f60c6827adab16f7795e6754bee8014dcdeb
MD5 de50ae3a34b0019721aea10e0f0c5d0f
BLAKE2b-256 e13ee1c86d7320e93e36a321850aaadb7e2f349fb41fbe7aaff0878fe301ff9a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 8ef3a9301f07c8e845f3765458bc0309e70cd5d74397470a9868c4d5de4ca8e2
MD5 3930575ea897e9cbbb7af22263876a8e
BLAKE2b-256 b1b1cc47921454e43a1ea7a7ca9e4e2ee40efea50b10637e193ce6e14d6f4ff7

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 afe71643d9c9327f83e87ee90df36904b028464cbee9bae9ec62632ae6c98dc8
MD5 78d7c5ef96e01e20c210a4afa84b8166
BLAKE2b-256 f59b0fbc1637cac1dc8d1740d9b6d9bccaee6204a3239a375a1d4b1a2807b9c2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 20ecc8e553c73a8eb7646f7790ca034d2007e3f3c314c0add05ce63abe1c4005
MD5 78227930745a80682b93d6dff6395089
BLAKE2b-256 d70841359b502604dd2589b5610a75df47ff6bfd9783b4735399ec3bd984b09a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 3d0c1b1f4a0791a30e0528d40ceb9387f6b16a4d5209d54c2aa6d5ad0e9cc0e1
MD5 05d01820452ad09a87ce95d0fe92372e
BLAKE2b-256 7501ced985db377b2ae89db15bbf6b8335f73907ed74397e6a69e0106339dd9a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 ce4913b1cfaeedfec0c941662914c24b3c4dc6a30cf12f24ff53d522e37a3566
MD5 26b3f8a6b7e3b9c16aaf0c7dccf22062
BLAKE2b-256 1c8c55f5c85af76778daf491ab2308bffdd3dcb4ae03b166f22a4535c855f9bd

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 65b7baac6ac8829d53a0df8b8f4f52bff6842f117524187fcdba38d0fef79320
MD5 46d481e18a4fb0e7ef82daace319f529
BLAKE2b-256 1c0ecb22ed54aad06c19cd281e561b75b986f3dafe7d3693d559df2ee55b3974

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 13a1e74ff5f5ccae18ed5cb279c552c731bd018a77172caba227abab486a2002
MD5 ee6434b9ee6649dd0044cfe1d3fb2af2
BLAKE2b-256 f0dce035280e0990aaef5c97aec0379db958f6ae69c24bd58cd01baaed906183

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 e2c9facebb5f74591f264230d01c8e5ac6f57f725395dc64ffe044d19c340c76
MD5 342f161ae0276a96ceb0b2b9825cd125
BLAKE2b-256 ad24179a9b5150560b1d0caf13f55c10895ca8809e1e7463423e47f69f7a1e73

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 544d6b49b126fea06624ec36302f3796501adf59b19d581a7e82b57333c66db9
MD5 019be55101c1151318d0a2a955009248
BLAKE2b-256 eece2bd928187b2132c42ae952cb4448b76978b86259bb1bc2855a90a8618f64

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 2d9bb636e1a69803faa8d07a8a4428ed18a5a3ad870587ba94deda4f5644ce26
MD5 8d4b97f73b48e8382505519768ac7c8d
BLAKE2b-256 ed875d8cf18449f73e36cba9371f23736263e954914820587bab8f005afc0864

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 574dceddb28d78f18a0f4dc9f802960c4e22a2c7e8715f4ce9a8dfac21d4a909
MD5 6bfcf1c1396180b099a996469e9c0dba
BLAKE2b-256 a4b338c5d5ab762dc9ab3d28eb780485925f27181802c2888cc26b53ac61e9bc

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 35076a6b6678d7bab555832a8ef44be9394642a885bc47a0051868364e0f2488
MD5 5290763b86582ecb34057e13c7332084
BLAKE2b-256 3e40794dd59315d9556c3e1c59a99bc08bfd9d8adeae8be417954c1c28e3885c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 7702799ae0ddbb36dba322d7c1470d49b8b7ff8adb3737a97f7e09ff68ab766a
MD5 fe0dc90134ed136e38daf93a1d802b86
BLAKE2b-256 b9a955a24af1d55fcf5d9aa66e21f91203ee50fae52f4f2fe10c44b6fc9ef6eb

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 6476e636fc12cf0cdb41a6bbaf453c17b816c3a35f313415fabb0cf7a50c8e1e
MD5 0bbb0f3fed6e92ce924514b9c561541e
BLAKE2b-256 7e05a832af1d52cce9abf586e6c18d0e8aa094aca96d94871da64b29a3cdbdaf

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 ed18f5e1b396d787e43381fabdd9e0901f26745b1ab51b8127df17f8a2d56f84
MD5 721e905d8eafaf7459b04a3d7111fb58
BLAKE2b-256 8f36313097b2a24a03dcf50a1d53b1dc2733e41df6f12bb3705f8d174c174026

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 b0d92f4a11dc534efe93f688a42852db8e4a3a8c8bb8c11a2fd5b34c489d7b6e
MD5 cfe6a515cc5544a002d7120a3f2834ef
BLAKE2b-256 daf1775111342c8e9a59c13718050a87750d86a71b70ec73afb9949469b9391a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 884671054654b2b9f4fd7bc5d3217205a8b7c60ade6de5e404aeffe95b360747
MD5 b961e6b39acf2516cb23a65da50eea84
BLAKE2b-256 ff353c2c494914ff05522221d923c85c69bb460eb0c5200fa2ee45fb565c915d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 8943289189bb609cdc92bbc6f979fc4d02c858cba6e4511fca1fc30985c90c12
MD5 3826d5cba559eb3902ce3d5ae40af36a
BLAKE2b-256 592f4b4f25a74675e1ddc9570cef0dc1c4789ea4255a8aa3be4c1a2506c750a5

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.7-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 1cf788bc030678ddc7761eacd60666d5983a47a64ed52c1f200f0a43df2b31ce
MD5 9aca01875132a9bbd908aab5c8d8767f
BLAKE2b-256 f8c2099bad7c524a4cabaedb821c7bdff295a7e2953a5a06f51cb6f7426d7530

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