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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.14-cp314-cp314-win_amd64.whl (650.8 kB view details)

Uploaded CPython 3.14Windows x86-64

tensordict_nightly-2026.8.14-cp314-cp314-manylinux_2_28_aarch64.whl (588.8 kB view details)

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.14-cp314-cp314-macosx_11_0_universal2.whl (577.9 kB view details)

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

tensordict_nightly-2026.8.14-cp313-cp313-win_amd64.whl (648.8 kB view details)

Uploaded CPython 3.13Windows x86-64

tensordict_nightly-2026.8.14-cp313-cp313-manylinux_2_28_aarch64.whl (588.0 kB view details)

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.14-cp313-cp313-macosx_11_0_universal2.whl (577.8 kB view details)

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

tensordict_nightly-2026.8.14-cp312-cp312-win_amd64.whl (648.8 kB view details)

Uploaded CPython 3.12Windows x86-64

tensordict_nightly-2026.8.14-cp312-cp312-manylinux_2_28_aarch64.whl (587.8 kB view details)

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.14-cp312-cp312-macosx_11_0_universal2.whl (577.7 kB view details)

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

tensordict_nightly-2026.8.14-cp311-cp311-win_amd64.whl (647.7 kB view details)

Uploaded CPython 3.11Windows x86-64

tensordict_nightly-2026.8.14-cp311-cp311-manylinux_2_28_aarch64.whl (588.2 kB view details)

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.14-cp311-cp311-macosx_11_0_universal2.whl (577.0 kB view details)

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

tensordict_nightly-2026.8.14-cp310-cp310-win_amd64.whl (645.4 kB view details)

Uploaded CPython 3.10Windows x86-64

tensordict_nightly-2026.8.14-cp310-cp310-manylinux_2_28_aarch64.whl (586.7 kB view details)

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.14-cp310-cp310-macosx_11_0_universal2.whl (575.2 kB view details)

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

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 b932270b12d10fa44152efe4001b3a024b2077a1b7b8b9bac7e53368ad7e0048
MD5 f41be28869e8f087e9862e912af951af
BLAKE2b-256 78d56ef63d320784dddcfbb39f2857bb969ead9e79e056948aa07364e257ce82

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 37e88202e619cef70cfcde6523f6b8e9c7f6d439aac2fcc1653f08b75166caa7
MD5 2ecb13fb370f76c3c4a0802bf11818e5
BLAKE2b-256 175b1f8db50a29825d4f6c2be0e4cea8974b64799c02afa040fd4b4c109f94ae

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 4cbf5cc34f9f85215937125f9d3fe62a91bcf371066184f4cf1d0f00ee944633
MD5 48a7bac7d06c7303ff27ef3463284ca1
BLAKE2b-256 2203cd7633d7326364287bc0042b691125cc5c387d784e922fac1ed7b1e7d520

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 06d354dafc5368baa2b137c175e44cabc09408393b69534b2b0d4e61edc6cfc1
MD5 35ca0513bf153804232c7b768ce7a3c9
BLAKE2b-256 034d366b36e5551197268f2cb5de138e78f0986c3cad9804f9b67ac8e6c2a9cb

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 94daecb0d66fe27ee7d550490a65732d6c02c918542dbc8f76c785d7522e12e1
MD5 7dce1a0011c3603b82912fdf6631f899
BLAKE2b-256 c392e07a5fe40f9a1b6e3df7a141129843d6b64f5fd17429866ca107b49ddb41

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 11e169bef200399df7a713183b6f68af39aa6ad6437762bfba0a9034349115c4
MD5 37c99443ef11f57f061265aa03f1a698
BLAKE2b-256 a046a2301af844fb9d9eb7a2d58f7f76f9983425e19bf12d801053833a5da6e9

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 0080dc71755529c004278c6120a3a532d5d753ee195d339ab558aa969d084e39
MD5 f26aedf6474c9168e7ac39cb67ecdea3
BLAKE2b-256 59354eb0f35f23a763797b48fce7abff036a3fb2b014847d9e2b77122b86c1b1

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 caf0895661396bfc82fe008c6b6af138bfdb3a965f6da44ebcf2b5a3e7014e66
MD5 4cc9adcca6102e880086214932d389b5
BLAKE2b-256 a852391c4a4185d7fbebf2e44dc22e89cc1468fed523eca296f63a73b08f6d79

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 c5657f673375b5f8d4c7926e78c5d4f23c1a6764571252614d52c5deeb0feb5d
MD5 694c377cdfc882ca51d6ccb34cf59725
BLAKE2b-256 a06c412d9af940c9de226830c0f6279d3c35a98c40b0ca92d69135e2c296ba51

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 12d81aeb3c85f1c7829708a3236499b1f390c5e7ac482743547aad201f27a47e
MD5 51e4b0a7964305752df081c6c0bf4585
BLAKE2b-256 e0d82dc8f60e4b9a6df06751988cc74f3570d6fae0913ec098b4ea208968879f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 156de722b1afacd67708c6df2f57cb245304bfd3b5424ffc70cc6111fd707476
MD5 157ab51acf29b15562d9dc2a8faac7d6
BLAKE2b-256 822597d85c6580169a1698606d7aaddb9111d1b144fd6c0a459d2c40fa3885f6

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 2a20dfc04754847157cec8429445f23531c6e9e36b4c5eb13d3678e1676bf644
MD5 fae1d35663bc9509728d1ea013420557
BLAKE2b-256 9c1a827cae5365a572cb345eb8cf61324a56e77ce63c47beab3e966721cc4ce0

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 1ac02b3c5dba8722a795a6730bc0a2a09fa759df6d915039d62dbebd823166a7
MD5 c96140380d03a4326b96c4032419a1c3
BLAKE2b-256 a85c4a30db25654a1d381f156c54569001802ebc0d2967a63c8d0278130e7227

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 44b50fd39fa2e076e4d0a3b8aa66396929ced90b122cdc32c957b067392d222f
MD5 acfab23ddc7a6598f04215fdda02b124
BLAKE2b-256 0d9ef6c0dab4dca1b36e24ef415f414b3791a970a4b666d4ba44eda44016bf5f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 b3a1f3c777b2d26c0ed19f88276db18db6877ff3c5b3762ef4945feabfd95e2b
MD5 331fc5e07287af33abbdaa2274f20118
BLAKE2b-256 22d6b99dd76f232b7b04f11ad4d276ff2d7eaddc422196fc75bc5806f3d5aec5

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 7a772215dd38cd7ccffcaac1d4fc59d4838810bf1f2e5b9ec062d6b84ee26d20
MD5 be1a1f16ba9d40a6a04ae72e47906c0c
BLAKE2b-256 f4e6311090a1d3070050fe289b0326b2557b99aaa2874566d7ad30353b5a517d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 7f5ca040415002a541d0f827b73a915f3c11659f4ad19264ed8064f43bfb170f
MD5 bf65336a38d835053dd3c5b5d63dc4dd
BLAKE2b-256 a3da91b0c1ba6fd9530369bf1b5bdd2c7f417e3c5b731b21af5927ad0ae9a9fc

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 7c25fcce25501400df8ca59fa4d9b489817c730b4f7a21406b89a2f9e75702bc
MD5 28877e26939226628bbde53a13a3d450
BLAKE2b-256 04f5e046813464d8fd3be2c8ac66d562e232210805ac1c9642290df579155069

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 a81bf29b66b5e02702f0e025f2da274b3537586d29a9c3d9155f36db88c40f85
MD5 c20ba8905174b4a182ddf95a8a095285
BLAKE2b-256 81fd849a0f05cdc686c2646ad02d42a65cfa7526d0dd870c977bce020b60bfda

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 93389547d9e4262c933c314cfe33ffcd10e451f85e27bca80c7e9e470bd1fc20
MD5 26531008868b1747b5f1a0b3b5f1e971
BLAKE2b-256 eee97b22045db00e3d9222af6b3b81144a7a75ffa54d23ef8afdc789b381c185

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.14-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 f6ae8a161b26af06609e22b1ae5cc23e2f4e401dc8cdf0b50e681c6284e7d440
MD5 96e67ee767f8ac67900b60a3fbf303de
BLAKE2b-256 18a60987691e27bffaf5bac15dc9e9bd501fba07352eb1f509e04c89444f586b

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