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

Uploaded CPython 3.14tmanylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.14Windows x86-64

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

Uploaded CPython 3.14manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.13Windows x86-64

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

Uploaded CPython 3.13manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.12Windows x86-64

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

Uploaded CPython 3.12manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.11Windows x86-64

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

Uploaded CPython 3.11manylinux: glibc 2.28+ ARM64

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

Uploaded CPython 3.10Windows x86-64

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

Uploaded CPython 3.10manylinux: glibc 2.28+ ARM64

tensordict_nightly-2026.8.12-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.12-cp314-cp314t-manylinux_2_28_aarch64.whl.

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp314-cp314t-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 84518c8547aeb3b0d9530b803a2f15c23999d0688b23c3d5ec592bd3475dde47
MD5 52f63200a81143f15513276429d24f7b
BLAKE2b-256 325516d0aea5fbab0944a8d4a459c32357dc1f600a458b5b88a268579389b5a2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp314-cp314-win_amd64.whl
Algorithm Hash digest
SHA256 aa10ab457bbb2821738c0193fb244980adfbeb78ddfb394ded1f4cb290621e44
MD5 1b0aef17eef3afd3cbd1637b971caff7
BLAKE2b-256 773e0b966309f6370d7d28be6c161deb0ad92fc7a946f4b6784c7520b6e0e77c

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp314-cp314-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 303bb1c533de79239e6ba3b8ad3e536b96105f81d9dd23095e2131a4cd035009
MD5 e7394f66256dbc679f1acca44ff03ec2
BLAKE2b-256 34b958f1d7e52b83f10573bd3ace996067f4e92eb5f1a9f4e0cce89fb2a7e8b7

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp314-cp314-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 454b8d59e6599245b3ef7ba17e68f247c34e1b212d5065c9f9802ba5df5b24d7
MD5 e219af70809e8268a378a23d95b03e13
BLAKE2b-256 a3ee50215e6ff4b5f01768344a96eb0408c40a2ad7fe80589de093649db8ccb5

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp314-cp314-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 86d3debd21a2e5aed263e02e1144ffe808307297c00cd7fbf90caef17f71a8a2
MD5 1047a6f5ff239053f19a8a356df4fce9
BLAKE2b-256 d9ea4996e8676a5d72608fae6746bebb6b18067057c16481302fed963b599ad6

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp313-cp313-win_amd64.whl
Algorithm Hash digest
SHA256 b25200a5f16b70396925455affce7fe6496342959914d4c5fb5d76225071f909
MD5 f61832f60e560a6ea1629ce822fd5a21
BLAKE2b-256 3d4781c28c1f1038798ce950049ca82a820ba9657ffaa8dc3dea71d6d96f7055

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp313-cp313-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 9db076f8e3b0fd347183b35cf631745878bc1d41e7059705c13b2840505492ca
MD5 edb10a021bc44f2f568f386f2bcd8fd8
BLAKE2b-256 09fe347f7911437fdfbf83a941b9c8af547b64fe4bf937ab8a1b2fce3112a85f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp313-cp313-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 8801b8e82d495cbf24cfafa983bd1838f3baa6d9bdc0c3215e790e83fd3b0741
MD5 effc89467c61b7a81a9d22376f474ac8
BLAKE2b-256 aab54168c683fb0a82842a41a1e79eed588ec35fab34ac439919a30344675bf8

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp313-cp313-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 4659fba13567e9229923a5197af835979e6135a3162b824c27dba1615b5e1247
MD5 fd0c343b802d37ed9577b34577cf43d5
BLAKE2b-256 beb7c06a9554ecd05030f23d6c44d70a12cad0c877e925161b5d78766c710368

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp312-cp312-win_amd64.whl
Algorithm Hash digest
SHA256 f31cb7046a04b99a1ef6ab4ef4cb5ac6204982863ec5743804408cd802797aae
MD5 37f63daa957c850cf964541d9ae4c89f
BLAKE2b-256 d1db46fbbe759b9e97377492a5957aa342146ad08dab5a65e9d4fbc06e8bcae2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp312-cp312-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 634100c67eca030dcdcabcc38f642e23ce461c42bea9b68e4265b0e916ae00b0
MD5 71adb7325d7f1b4a533c0bd4c710dc36
BLAKE2b-256 b57520aadedc29a5bab96a863012349310aebf38bf4e7ac8e8f807e3a8ed8d2b

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp312-cp312-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 4bff99165892847fbbe89998a8cb9fa25891de205f63f7ab48030cc7dc9c0ff9
MD5 91c888ba58f2ff208d02f777ce21831d
BLAKE2b-256 e6753d3bf6f32acf9174b4d8d3eb343a62badd2646938b420d26a3c75e6f914f

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp312-cp312-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 fa969fc2a9bdc41e7628807ca5ec37730cc3b9004f91cdb42a2fa920763f34ed
MD5 fb73e1a5b219a409ff092bb0da43b42e
BLAKE2b-256 5e8ce49a7fd78ba29c6c9e6d0b67c73c0c691b93d95c09d084f4b23709131ad8

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp311-cp311-win_amd64.whl
Algorithm Hash digest
SHA256 21f319b6c12e3a24f9eee5d606b6c177b38d6d59e4cea7c47c3e4c7a9f621508
MD5 82f63ba72b8d20329c94daad7562eaa1
BLAKE2b-256 b3334053869da1b664f3264e5d5cff44deaf1bec87d69ceb8c4e59c1e0dea4f2

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp311-cp311-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 38dfe5581dbb8ded60ac7772ec10ac9155ff52397384e93381a7b24ef3563de9
MD5 583aab7e97f7584268095345ac398637
BLAKE2b-256 e404f2e7a584a97772e08cc082900fd95e1a83583f3ccf92d8f2becd93a59b97

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp311-cp311-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 02608700ddc623e75b9be9ebd25f3abc66f0596a976ef621b5d0366fcd8d9c23
MD5 168f32c151bb1f3e3672f71d987ccc8c
BLAKE2b-256 aa67dbf98b22a3e37b10e34cf74763014e34bd4477fe5af51c4227a23371310a

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp311-cp311-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 038effca07131af3238bd139964569a23334871f7c5e105e1c1913bdeb0a2f09
MD5 53d1539520fc9a7fe421bb47449d848b
BLAKE2b-256 0f87831edda2819952023e39cb2fd6e0cc5eb67dee9fec817b67d956f668355d

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp310-cp310-win_amd64.whl
Algorithm Hash digest
SHA256 9d13643e99b1c7b268822e7928c4e7c65ab7e8791361d6f9e8d24b8fc42b98de
MD5 35bbf564ea9fcd51ca7d950db8ce0559
BLAKE2b-256 22b2ebd566907bdde4da8b0f0420982c08c99759979391f1d0a3cbbb4ac349ea

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp310-cp310-manylinux_2_28_aarch64.whl
Algorithm Hash digest
SHA256 9de33eed07b993526647f71e0f8df8ae79ad520fac34ffefda86c21348dc0d64
MD5 6867f9d589a512d84bec78b0b3e1d6b3
BLAKE2b-256 5ecc9785aefe70c011bb2b604556637be380e9372468ed1d0b2bd1f83fd74326

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp310-cp310-manylinux1_x86_64.whl
Algorithm Hash digest
SHA256 4e6df461be36f2a9fce736073a8a484b9eb8274150f576d794052ddba915aaa5
MD5 b94fc6a6a131edc73effaec13a8cca12
BLAKE2b-256 cd4e632a73065c3e74afc50c837d07f0d0654bce8b4c1ae3e4e3c31e2c6d6683

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for tensordict_nightly-2026.8.12-cp310-cp310-macosx_11_0_universal2.whl
Algorithm Hash digest
SHA256 a922e0e28bccc29c37833769e31b014f10090960e70f45b1b51ca6e017a0d8a4
MD5 bb7d830a43fffbc1a0a24c1b743bd4aa
BLAKE2b-256 db0c7869ac44ce731daa0b243819733e463b828a0efd564891ad98e132896ded

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