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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

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

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.15-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.15-cp313-cp313-win_amd64.whl (649.7 kB view details)

Uploaded CPython 3.13Windows x86-64

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

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.15-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.15-cp312-cp312-win_amd64.whl (649.6 kB view details)

Uploaded CPython 3.12Windows x86-64

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

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.15-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.15-cp311-cp311-win_amd64.whl (648.6 kB view details)

Uploaded CPython 3.11Windows x86-64

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

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.15-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.15-cp310-cp310-win_amd64.whl (646.2 kB view details)

Uploaded CPython 3.10Windows x86-64

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

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 ad254f537c870f9a65b88b9249fa2bea02d40731d648ac1f1226befb24f4b720
MD5 162e88f301c35704f3fb4b649e6a1b21
BLAKE2b-256 f9e4a5958d2fc157b0def3296d71ab1d815ffffebf6a24272ea2125daac41f71

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 871d1a4cb9cdda77340d55c104d498be82eb30a8d7016dac4ed4796f8ea91747
MD5 5de82117aa475ecd17ddb314b6f6f150
BLAKE2b-256 d241d14286688a96b6fa0acea4f32f2be6bf5b897608b9f141d552de22a0e17f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 8c54f5b55dddde7ec328b9b42918b2111ca3eb8db3d7fa3de7d9be84ce8761e2
MD5 26f9dcd8a4b7882e6beb310440d4b810
BLAKE2b-256 7bf5c88e9797233372beb3681cfc69029cf46c1f80995bb5f80a22049ded6f2a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 0cf4325dcacf0be638f15c1015e42d691b302e9d375062d0785629f322261e45
MD5 1fea66935384ce6df228f2d48dd86e8c
BLAKE2b-256 7b687835e3408df95246a8c5419067fa163ded80fbb63044278f5d1284d90b22

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 47717057ecd71e840c936acd23bf8589ccecebe1dfd54c67aa57286c303bad62
MD5 70dd91a8088324d5b08a5337a627d7f6
BLAKE2b-256 e9ad6b59bf80dcefe5a73e7afce2f608cfc0c02527d7f63f2b9c8d43cd79abae

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 885972607d9ed58de663dd22854f9079aeb7778975efac5c1fec2a505a981f9d
MD5 83e4952fd69c747edbac2bdaddb875fa
BLAKE2b-256 2a4fad839cf9fcb6b49a01fdba16cb24a27355e7691be5f91f8ed2d4c592770f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 500f058be24910a0ae32129c1986fa50807675052dac06e3510087dae91c8502
MD5 3cdf439a0e9cd86083c3afc4e0763b18
BLAKE2b-256 ddcb183912976bb12f9bfcb6259d1d1fcb991ed8a21a75624941c847e11773c6

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 a6b275a8fe7e744125e843363b5ef5d0f7091b7afb41ba9c476968438121484f
MD5 cd2e2907ce85fc0436f2b22943758093
BLAKE2b-256 97c93bf01796ed45cf6d86e14cf6450e9271952171da70cd4ac5aec22f7a40cd

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 f62ec0d8906dca9c0850466b1160bbd8393b8601d12d44050e4290703a4de1f4
MD5 301d93862aac9ea447f5c5df2d72a18a
BLAKE2b-256 1bd92862f98c94c2824e01ce7db6b76022c24e200b97ff801e01eca1cc33b660

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 fd372d1a0e70ef5027d67b5ad70d9243724ea986287e47d7c16dc863c7e0339e
MD5 019ebbf5cbb3c95af48c3cbe38d983e9
BLAKE2b-256 ef1de8bb931cfa8cae65ab5505c093032fae6bea90c5497b8f568bfd375c0c90

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 df8ed95063eeb57c13644fde644def1a8a82deded64b75d0e7498f383772bf5b
MD5 c18069a762c73afe8314f7258e1735c1
BLAKE2b-256 6b18dedb15a18b5af0e3fad9bf776c933ccb54414afd196bd8371fbe9fa3307a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 9339935126a29e1281488cf5a7bebbc96cf49761a2e446b872122ed570f67c4f
MD5 2634331597068931652e3cc12d2f29f6
BLAKE2b-256 cbf76fd2c2e8d6fcc13034ed3e59258ba86e57f2179581d5b08f6e68fa9f70c1

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 86ba31d385a8b67f1e63b8e74930673a94a3625824b194bc43706abdda9af9b5
MD5 03b520a8593674328d3ff70ad86c667a
BLAKE2b-256 51483cbc28d2f21c50da9df7f6a03271bdd13584f4f807761c5fc7f458ae42e9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 bacf07722395beba98319bc011f93b350efd2b33f4c06fdeca0955717e76efaa
MD5 3ec966e55b921f5a10b6fe00be84a468
BLAKE2b-256 97051a7fae54f9365c3346d7002841522b687fbf862d6a6bb523fbad3fea69d7

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 e404ad40838080879b4209b7723dc2298aea368b2e7f2ac9a9a8c2507d53e54f
MD5 54e1819da0fcb1421ed867bc92266039
BLAKE2b-256 eef3e7d2919b451663a753fecccee1a0090ab4970c8a4aaeecd2bc00c473cc03

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 8ad8bee1ca49767c23800df5900d3723eb63e8b5c0a238726fa52c8f9e8bf823
MD5 bfa4c616f1b572cd1da799dc14027916
BLAKE2b-256 b9eab2a7ef454fdb75e6ffbbf6564063475ed69095e7358a0224dbf9309f9423

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 f61dbd09db123d485df26851d002c1962f2173b592a4111bccdb6a54d757ec34
MD5 5f5d4be3b5c8b3a0eefe848a24b571ae
BLAKE2b-256 0f0a7c41288dd8abd7217d6a50c4d491aa849867294e4f11afca1330a2f31690

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 0c1fa7e6b23b6665a1fa64490f376929d44c6d487e390231783f4e57be902bc8
MD5 e64a7256b3f73d688c659633e4150a4d
BLAKE2b-256 0b888f2b25e41c8ffba843023327bf363d928dfb83c20b5e76e7c5fc41b9bef9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 2e405b1fc2dbb8638b0d64815f9cc896dee308fd1020d79111feeca3438b9e5e
MD5 5fab15cea2998a904e577d7d21ac1ade
BLAKE2b-256 0f52b558c73033adacf102e557572a5471ec984da24a1152a4a2de3329937566

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 304c3fc74b8ea0f3586f7599d3003d2f3a7e9c48201b96b390077a5b960f0651
MD5 a7c57ff9889a8d9015514620baadc2a7
BLAKE2b-256 a13b94bbf649a53392254d8be450ad1becb614d5deee125025686d09179bb59c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.15-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 3cfbd8924edaf28d85e95386ee45620fb2e3d58792bf75cf97c0dc954b95bfdc
MD5 eea4bd481575dc416053840eb41c25fb
BLAKE2b-256 7a6042a34d4173ac05008794ff16d199301188714a30b208de32ecec8cb7fe95

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