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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

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

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.13Windows x86-64

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

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.12Windows x86-64

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

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.18-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.18-cp311-cp311-win_amd64.whl (648.5 kB view details)

Uploaded CPython 3.11Windows x86-64

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

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.10Windows x86-64

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

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 9940921fdf02e39d4784899eb674bb563525b5778139787aa6d8d605033c01cb
MD5 85a875bdd3551283c8481247c12eb03b
BLAKE2b-256 72614686264f296bab03171cb1d9b2b874f924a01ea667724f7c4ece89e35e48

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 237aa48a9cde664fd71cf586dec48f2e9b496371c73b957555d665a448ff974f
MD5 f2b76854a9d344a1958594d69e3b22fc
BLAKE2b-256 7abed9ced32125dc998f09429e62a0af92714ed8999cd01a3b93203436687afe

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 ed11a55d757eb06350c91cef99f5050b04f16548232778d4e84708688c367862
MD5 88be015aa879e39649a45e6737fb1cef
BLAKE2b-256 24a69b453192184264064c1906770738ce7cc54e12da44ba3910445737a282d0

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 9c137850649b615463a02336e5ce064a818865657b3af41cef2708aa7ccd6133
MD5 29d9fda2c47e23337281b437d1442a5d
BLAKE2b-256 5f373e08e6ae41129e0b0cd199fc582cc1db64c8378772d70622fb287192388a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 0096a102b7d2f05a004c4d59cec79b048ea929baff45d702b43b93dd1efe6b55
MD5 0dc4cfa2f755ea0b6b17346be37b86d1
BLAKE2b-256 b83be5f61ed213a32edac6440831df89aac75bd0d78bb108e6ff5a058a68d24b

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 4bc109611542ca903be9d2a6adb3afa51f5d42678ee4d43ba611f2c969995f49
MD5 c38a0bc7427b1974c61ba049f2ac2ab8
BLAKE2b-256 255855e877c6292bd71265132c2d9390f25e14637ca9e8a069a7573398b40646

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 29cf3f24c2c8fd0fa07505766f58db84843734a98a3ed671fe798565b37db486
MD5 acfbd57b6fbb59f33bc04c35e786764a
BLAKE2b-256 b37c981b055e93a5cffd5ddcdbab2b4f8006f1ccdf5ed5ae9bb25a304bba658d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 a98f75534433a6a6c28d956bd007b4ed4115f280f757f66d7b883abb11e9c674
MD5 a4833c50e6b2b60839e495da8a5b3383
BLAKE2b-256 3912cedeef9dbd6d8e03c796d5812a5106df0d7c9fc7fb8964b36bc938ebeb4a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 f8888534a0fff375f1785745136362985e5a6ab3c564a249e9200aefbe3d7a1e
MD5 921cef31d1f29bf112e55b36f81d4c35
BLAKE2b-256 77c215171b498f084cbafa70f336b9ad0705ec5a8bb9ae6ce73eacff9b82c853

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 aa7f4f3baa9f5ff2a2bb30ea6a6c462b0e5ec771d04f23fa340e829caea99ca9
MD5 bba292af4f4d152d361b02ced4722196
BLAKE2b-256 79576aa18302074f6cba187e5e8a982e1fd7cab9276b958e10d76a0fe8987b93

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 4cd7730a4cfe676148b0d53bc6db4db9cf35cf848cb635a64e607e9d7e6930be
MD5 18cc3885f48aee9927c40ba3418bd66e
BLAKE2b-256 43d9f116c0e93da00fa8240b2a50d52b28258a0b71effe7d22f26968f65cfc63

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 52f697e26a45e67d427687b3bab0984f9509c82a5238dcd74f01af555dcb9ef9
MD5 66418aad9869da939d958479478e8703
BLAKE2b-256 bce799520e268a6012acfa1181847522dac66ab81e6891aeaf7683ccc1104e2d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 85c4ff5d672674680a4f5c2b89d37e2ba1339e9d760aa31940c4c43cdab290f7
MD5 ef932e0ed9a34fd594b0d535c7105849
BLAKE2b-256 aa9215a22a44da35118bba4d01ce874c2f49729e2fd0fe4feaf600735a389930

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 38bcbb3dbf8bea859dcf5c4a94c4b06b8eddf43ad06c3326b101f88e3f792b6e
MD5 afe9f2b4c9834b8cda26c2b84e5aa7eb
BLAKE2b-256 dd26bd1d1ccd4c6d0fbff996532874eea156d4ac9b0f10b70f98375da3e18789

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 ce0fd63b3f7d81e818b6277ee59798cfaf870b1c3851a4a8b343dd75774df6a7
MD5 ae41394c9e590d377df8c3648385604f
BLAKE2b-256 df4f076c8e0dd4dd74d7eddd292b499b6c284776dfc30c685cdc5ad62d1fedc7

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 518fce6648f6124e9a26727e090ef860614fc0b40810c15e57a9f767eaeab7e6
MD5 bd085fb852ad856ce9d67a0a583f3446
BLAKE2b-256 830e9d2c07cfa085adb0c1c1ff7586c55afdb85f3852a80750c905f5861c83d3

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 398a8d9667e183008e9a8fc204c2094d40d3fb340a2197c331b8ef199f4d0832
MD5 9c01c0fbadd1c3bd18078e19ef863310
BLAKE2b-256 3a3f748b5489d1a28252318793cca24da685d6c264df3d660e4d54bc17a84a4a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 219365cdf904b6227dae417a59b0777bfc65e7cb43c37a8c80852a58374cfe07
MD5 583d7c72d2c7eea8272e3eb44135dfdc
BLAKE2b-256 ae555671d897dfddd6a9b4ed23e5a67d4a780d242d2cfa3f531e175376b92c64

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 7a7531564ad773d958bb6d0ca4253014364e37955f724904d61bec6ed91c10ad
MD5 de6a63b34428aa19b40ffcfc5f1dbfbb
BLAKE2b-256 3dece4cef25924b0b19a86d8d8aa47ae1fbbcd21a2542aa7dd1578e12fc03f16

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 8a16e16acf50ca4abc159a1ae36781084e2647de23216fe5ae76d09d216b6483
MD5 81d9a303b3d06250b622c77da1051eac
BLAKE2b-256 41d33c5356a7efa98cf37d2c6ead30b7809b4de70371f76702f4521600050d81

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.18-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 321e8f1c660c921447433962721e43d2579245ca685651d084e2490f2b035f39
MD5 f03fcb8b45e270f47959e84ad1c3a8ba
BLAKE2b-256 ef1006a097cb18d947f3164294fb4883e206e2d092711c8546f66d499ddc0d6b

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