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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

tensordict_nightly-2026.8.6-cp314-cp314-manylinux_2_28_aarch64.whl (586.1 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.6-cp314-cp314-macosx_11_0_universal2.whl (574.9 kB view details)

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

tensordict_nightly-2026.8.6-cp313-cp313-win_amd64.whl (646.5 kB view details)

Uploaded CPython 3.13Windows x86-64

tensordict_nightly-2026.8.6-cp313-cp313-manylinux_2_28_aarch64.whl (585.4 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.6-cp313-cp313-macosx_11_0_universal2.whl (574.8 kB view details)

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

tensordict_nightly-2026.8.6-cp312-cp312-win_amd64.whl (646.4 kB view details)

Uploaded CPython 3.12Windows x86-64

tensordict_nightly-2026.8.6-cp312-cp312-manylinux_2_28_aarch64.whl (585.2 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.6-cp312-cp312-macosx_11_0_universal2.whl (574.8 kB view details)

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

tensordict_nightly-2026.8.6-cp311-cp311-win_amd64.whl (644.8 kB view details)

Uploaded CPython 3.11Windows x86-64

tensordict_nightly-2026.8.6-cp311-cp311-manylinux_2_28_aarch64.whl (584.4 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.6-cp311-cp311-macosx_11_0_universal2.whl (573.9 kB view details)

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

tensordict_nightly-2026.8.6-cp310-cp310-win_amd64.whl (641.8 kB view details)

Uploaded CPython 3.10Windows x86-64

tensordict_nightly-2026.8.6-cp310-cp310-manylinux_2_28_aarch64.whl (583.1 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.6-cp310-cp310-macosx_11_0_universal2.whl (572.0 kB view details)

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

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 718880529d3effcb1e41adc31fbb911b5f01919c548c3218f4c44b6ebf082c54
MD5 6913cdcbf3673638c6e40a5a090f247e
BLAKE2b-256 57f1f180bdcc04601a0990d749dcb19aa55c938cb383a2ccc2fd29f39317e3af

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 3cd202d02bd076d270832f453efaf84f442a4dbdfc438a39528de7dc12838155
MD5 b4f4d60aa84b1e9feffdde67448a2551
BLAKE2b-256 a9be8111571b4f5aac48a81c4f182cb07017c4ea4b2495a5f474d6c7403ad944

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 851bdcbf5f5003c4bf7a9b65909b2b20a9d2ec898c3e7c0cbb7b34fb35240022
MD5 4c21b20f522b9e5749dd42f754558b61
BLAKE2b-256 4750e68db58ec93cd591678f23df83b9d1bd0329d11aa02092cf886510987ec9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 da556f4df6b03d18729355cf60f6c49ab6afbbaa0f5db67109d7c9616b15c4e4
MD5 187d867c4733e9dd03520dc4344407d5
BLAKE2b-256 c0451f3cc64c7a853fc107e94ca484f31b36d0c4ec391fdb3d4d4228ac3c2305

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 8e85992daac12a51758dab1e7d43cb1e5867247a0f8af4a78d305dda06014536
MD5 f3372f431a50f92095e069663b8cc3b6
BLAKE2b-256 5e874a15e5bda664f1b8a32c3c3c1b08ecf47783cc7309026c431eb14a233e11

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 7bd7eb2b9cc95f1304c12bbb5dadd191e16e9cebfdf9c69d40a56d1dc622b963
MD5 512c0efb84faa802d809be8c0c2c44b6
BLAKE2b-256 0bcd1de112a0e9bd61198355a90f682d041928bd80276651800dba32be843307

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 524c59e56a09292ee902f60b8f6c01529276d830a75e46c93f3b809dc888aa0a
MD5 86442cb1c584ebe86e57a97472709475
BLAKE2b-256 fcf58a5d564853504ff565503f344b5e717fd985f353fae0048136ae538fb309

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 cfda440e6f210a8e3b87899ecad43c503e86339287968cb36ddcd69e88a25d50
MD5 373f618b206d8b6159df4c171401b004
BLAKE2b-256 3e35e93251a4984d4f3640907e48076e03cd5dada557946d2796a3242a3aacff

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 b5593d9b27babea45af386446c1f842a7bcc7cdd30c0e1c9a33fb23209117a35
MD5 bac4fe7fd363430b3445334175d0265c
BLAKE2b-256 8f7e414591c3e4554154a16595844945decacd53609e9dbaa569f771a69c20a1

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 fe581875e1af7a930f0bb02e4ffd024500a56f1b637dbc55827e61a92c6fccbc
MD5 bd7e991d6c8d8d1d7c943965483499e6
BLAKE2b-256 8abcd0288705f47b71999bbf87e276104a2c58f41adad9048ef1ce311e5426df

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 3560bade06788ce3675debcc50d2abaf7cf5ad8b87092c5323cfdd947284555f
MD5 291745b7947c3e9b2f089e7b6995ab50
BLAKE2b-256 06572afcae09e7a433303c3a3cc56904f09de846468840dc1b4cc9f09e59b46b

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 634e98375dbcac55b871e00557d4b70cc0cb6217e6214f23ceb00ea395406a6f
MD5 127fb1acabb04d0d1d830f54518cbf94
BLAKE2b-256 f37881e8e32c5f902cec216c98fb79420193a27bfb666aedcd8fcbb9d337d685

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 f7ba1f4be951a616ba7c4f153a32505873035df0b471496271cd24345e827380
MD5 57dd28405ae8a9a7fdc146f47edf5679
BLAKE2b-256 6e35c749b048810577e80da16c38f4a5c76fd180f0d6bef871cd0f550b2f7dd2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 1f7e6fec0d1cd310f18adcb99818aff958ba9f9ac9c65298aeec060399534c33
MD5 54172e59aff44e77e6c1c1e37785f2f9
BLAKE2b-256 3ae4447ca81624442e6ecbab124f78d459b7c5a58f308691b9d2beb1a2ac9ad3

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 a373463fb53bdca941a239e5189490ce3803b14b87dd73a82adf13b8a87d3676
MD5 c9b9b356cf5684d48bbb7e8ad58f4e05
BLAKE2b-256 580f8f0e7ad84b904d24128a43095604baabb271720b67291080ba8026adebbd

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 2c50a907b218f0391185aadcf21c30c9c507c9faa813239a1db0eb3112aa1af1
MD5 8cca24feefd885d3c1a78b1ee8d48916
BLAKE2b-256 d89a5c06cc4cab96c8a1c90470e2c56e247959e1c276b9795cfed55652a0c8ef

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 1a597ec6781f64e1dd8965efbbf1697728fa859e160b64d2cbb8964bb1929fd5
MD5 0361cc70ab01e34a18414bbf7a518797
BLAKE2b-256 2d40e33a1ef58281864282c85e21c6f50ed6ede1af586cdb3976c4195fb67f9f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 446a4dca75255d845f8e2d1efa2719c9552b55a1e687a37a571d7feda16655d9
MD5 37f344317aa97fba0afc18129e1d6685
BLAKE2b-256 011c00940e0328c5d51c4da343c48fa4c2b3f4f344d640e638946685a7a9ca20

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 56d43805ba2647a077a9d8796c3dc574f0d1c67b485e8704c55e2f6fea168e8d
MD5 969163d6921cd6019ed56c0837c8293c
BLAKE2b-256 774d07f6b31729a868fdef7c3759939fe471f92370a903a3c5f0915cb4b5d376

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 52ba7457cee76a195ba9a0182015d8945c7b468342f4211ffcff5d5f033da1ca
MD5 3b5cfc1d49b4ed278195985e6a9bc743
BLAKE2b-256 d0ac89b075ed5819b4b11684877712566fc122673795d0357d31c77005b88e6a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.6-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 2420e8da997eb86491a1bbc8187f3d48932ae76d2c41689c23b9d40d46f0eaab
MD5 466a01949f5fcb9e66cac560d6ea40c1
BLAKE2b-256 9355f983029ba08db3a53c57c443aaa7fee88ebd3dc5a4451472970cd6bdd9ee

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