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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

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

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.13Windows x86-64

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

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.12Windows x86-64

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

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.11Windows x86-64

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

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.10Windows x86-64

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

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 3a68ee30281da094ceefb1952a89a435be70b7de6313c640d11428d031390c5f
MD5 cd7e0e74dd96eacfe8b1caed53b79a61
BLAKE2b-256 878913f8ea8fa1baa5b45d0495e53bc3afd9d705bfaaa8d5a43481a60a6044fb

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 428aa071c3cd88abe6f413cbb0b924afe45c461084a99ffefaeeaf8168d67446
MD5 b7ce9df6714ebb56ddd904d2ee19cca7
BLAKE2b-256 7016f89bb815ef28d15fd90d8becfc16e53762476bf3e4fbeabd1e5ffe1905c4

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 de164a26bdd46685b572926b36b50c8e37ab368f4cb2c907faaf4572b8154928
MD5 90eee74a0776f142892a0393339d30dd
BLAKE2b-256 23bbfc1e55338563f75d556d7a557217eaf5ee7ee392b48490d6321ae4a8a6ff

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 349a82d4c64e447c05aea5ea24d7d52833c0050f6263008b023e6821dd761bf6
MD5 66de0b4b93f9dddc62b48008d425db0b
BLAKE2b-256 054a1abb3c0cc7c9ee869f4b6c62b366a98571b105e24962781f17206cc0203c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 4e33a1bfacca24699765b492957f3bc5504d2ce366e2bb5a3dff4af327b04d4c
MD5 5220f87b84bfba42cd5deedcd82952a1
BLAKE2b-256 816a80a6e04973f19e26b05176c5a991ee8d33447ae07ff2766a3d2ee5dd4cf1

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 e8686d023e4f6dbb775b46f4ea3a9e46ce7c1aae4c8c9f25330b6701bdf84404
MD5 fed75c0f680ce3f90691ee9560fbaba8
BLAKE2b-256 3071f891c0fe4e0cd765d0c3635e6eb7b2a94545e0f42731d1ea0ce5ebb63b15

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 8766a635818433485f3592631bae548c4e47c1f136586fadb1eff0e325b474b4
MD5 c4c91c12dd68f1745c1b73e80b1fb604
BLAKE2b-256 11e0f371fb8b61026a5e93dffa436ee9bedf2b673f9b73819e7bce3cc8e42c1c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 adedb59317be752ad6093a3e8e1fda245aa6880bc5130af990893fe80c3bdfb4
MD5 8cd84bcdc915a102cfebdec86754868a
BLAKE2b-256 71fff37698f403261b645059f1e625f6ac3648f9af51ff836297d4859549dfa6

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 032cc9ef5887ce49e5d5b89cde2783bec52dec7b1f9959f5d2befd375aae014e
MD5 531b329decdd2f8f9288ef9ca98100db
BLAKE2b-256 6b6d5d58a271186a38bcfebfa003f7c87e28017268f283a402ea672e469ac70a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 e36bc4bacba529f8fc35f4d4008c9cf43d26e72da59343a334ab9e56500d76c7
MD5 2684a38abc42ed5d197d96bbc0b25515
BLAKE2b-256 18574a2890f07679d8292e1a3bfd28fd5c447e56fa9e5b706aafaca1c111f510

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 fcf0c3f75f47cc16deb7ecebd3060662af2260c7f122812a73d1e8e9614f65b1
MD5 a76dbee86c53148c41286b8535567238
BLAKE2b-256 6b8a88e7f72b59458392859794e73d28dd7d74a07abfd881bedcd56ff116746a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 14125bbe8a141afb2dac5ba05a549f69700ad4f3b08b9cc50da34d578bfc2dde
MD5 da98996be63f834f08927d0c204f59b5
BLAKE2b-256 22c83d13db527559fdc4154361f875e37dbe62ce34e31524f4b5238400e0cb41

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 bc731988a381af280a6a08c7d7351766e8b691d1897e9aad8f3f04e8bb41ecc2
MD5 311117ea71603c66f6b75030268daad1
BLAKE2b-256 695235de06a87a9f92d4cb0ebe21e6c60e3da5d956abe5b9fade035d6a0d19b2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 d6a10262a45800db02225de65cc737857a4ae533592b5952ba5bf6442eeb3807
MD5 aa43203f53a74b1b4855d1a11e8c450a
BLAKE2b-256 826fd1bd4930c54fc2ddd5c36115adfd1cb168d6e5a9a3ad59f5cecee9613991

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 dbd6475d9e3832f30c9c658e5f5591bb42d961c1806bd170ff298c3b78a1082b
MD5 40dd5c185c8d2e5005ecd789299de22d
BLAKE2b-256 de39af04930a0c6157e679c9ee7b6987ffb5b68c680e8d1c579673f8b0af1a85

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 18f521a7076089eb52cd1bdd84c04277dc3579b86cc66a3380f426d485c90861
MD5 6ccd2d2c2abf094c21686e6bcd61a161
BLAKE2b-256 df2e52275146ed689b40e070b9cfde7c0266afb710ae8f3b2ed26f03669bad5a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 ebd287e9b9324d8a837ced7cd0585611e31d532af48c99435e2792fec9abcd13
MD5 008ce67110a098eadedaa38cbea64236
BLAKE2b-256 ab4071805bef192a53c019a856dc840b0893a1ee743951c2d78ce949e79f3a37

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 47a16ec241574908456511a117ee25269ec3a213561a53b928b33f23bc64a957
MD5 59b8063185459cdab6fbf6a970b265df
BLAKE2b-256 c5b6b863ed4e818bb8584d7a20a6584ac3883b244b1c380e23273200f19925e2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 7ba2cd26a223aa4fce32b7f3c6b9cab3e7275ccad50918fe3eac7423c2a4a8f3
MD5 0306924fcdbb297cd6ead4e0ae334ff0
BLAKE2b-256 4eaa8adada7b9ca2fee1f48dda2da6dc57ebd0e33af6140b1abdb8d4e1f13e40

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 5a7d92a30fabc42d462b3ca0ec58d8407c8b77c7a97d545c2d57d3bae177b0ed
MD5 e53171d752e8934d9a7064f47a4d22c8
BLAKE2b-256 0b9ff24174311285230e23c6daffb43d297729ebaa2c0b6ac2eb4631d3ac1069

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.10-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 975e122a4834ca6d47946a29d5cf04ce3b0312b7b746ebea50215816b8cb6807
MD5 1627a5e54edf9342cf206983b6ab56ee
BLAKE2b-256 ccb63cf612bc0ee71b5a3af843ce3bf01023b7c658742d4c888909044faa5025

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