tensordict

What it is and what it does

TensorDict is a PyTorch container that wraps nested dictionaries of tensors and treats them as a single batched unit. Instead of manually keeping track of batch dimensions across multiple tensors or repeating operations like `.to(device)` and slicing on each leaf, you work with one TensorDict object that synchronizes all its nested tensors. Every operation—indexing, reshaping, device transfer, arithmetic—applies uniformly to every leaf tensor while maintaining a shared batch_size that validates the structure.

It is built for data-heavy machine learning workflows where the unit of computation is not a single tensor but a collection of related tensors (e.g., observations, actions, rewards, and next-state observations in reinforcement learning, or parameter sets in ensemble training). TensorDict integrates with PyTorch's ecosystem, including torch.compile for performance-critical paths, and provides utilities for memory-mapped storage, lazy stacking, and functional parameter manipulation.

Use it for:

• Reinforcement learning: pass a single batch through environment, model, and loss without unpacking and repacking nested state.
• Offline dataset pipelines: store large replay buffers or training datasets with memory-mapped TensorDicts to reduce peak memory.
• Ensemble or multi-agent training: hold parameter sets or agent states as nested TensorDicts and vectorize operations over them.
• Model checkpointing: save and load structured intermediate state (activations, gradients, module parameters) in a unified format.
• Functional programming: extract module parameters into a TensorDict, swap them in and out, and compose with torch.vmap for vectorized inference.
• Data loading and preprocessing: unify heterogeneous tensor batches (images, metadata, labels) into one container that moves and reshapes atomically.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

TensorDict is a batched, nested dictionary container that behaves like a PyTorch tensor, allowing you to slice, reshape, move, and perform arithmetic on structured data while keeping all nested tensors synchronized.

Yes, if you work with structured tensor data in PyTorch (RL, multi-agent systems, ensemble training, or complex batched workflows). TensorDict eliminates boilerplate for keeping nested tensors synchronized and integrates well with modern PyTorch (torch.compile, vmap). Install friction is moderate due to torch dependency, but precompiled wheels and active maintenance make it straightforward. No known vulnerabilities. BSD license poses no restrictions.

Install

tensordict on PyPI

pip install tensordict

uv add tensordict

poetry add tensordict

Installing tensordict

Before you install

Medium install friction: precompiled wheels available for Python 3.10–3.13 across macOS, Linux, and Windows, but depends on torch and six other runtime packages. Active maintenance with a recent release (71 days old) and steady repository activity.

License in practice

BSD license is permissive; you may use, modify, and distribute this package with minimal restrictions, making it suitable for both open-source and commercial projects.

Quickstart

pip install tensordict

import torch
from tensordict import TensorDict

batch = TensorDict(
    {"obs": torch.randn(32, 128), "action": torch.randint(0, 4, (32,))},
    batch_size=[32],
)
mini = batch[:8]  # slices all leaves
on_device = batch.to("cuda" if torch.cuda.is_available() else "cpu")

Requires PyTorch (torch) to be installed; Python 3.10 or later.

Verify before relying

• Whether memmap and lazy-stack features are production-ready or still experimental in 0.13.
• Performance overhead compared to hand-rolled tensor operations in typical RL or training loops.
• Compatibility with torch.compile for all documented use cases or only a subset.

Package facts

Evidence: tensordict-0.13.0-cp310-cp310-macosx_14_0_arm64.whl; tensordict-0.13.0-cp310-cp310-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp310-cp310-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp310-cp310-win_amd64.whl; tensordict-0.13.0-cp311-cp311-macosx_14_0_arm64.whl; tensordict-0.13.0-cp311-cp311-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp311-cp311-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp311-cp311-win_amd64.whl; tensordict-0.13.0-cp312-cp312-macosx_14_0_arm64.whl; tensordict-0.13.0-cp312-cp312-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp312-cp312-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp312-cp312-win_amd64.whl; tensordict-0.13.0-cp313-cp313-macosx_14_0_arm64.whl; tensordict-0.13.0-cp313-cp313-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp313-cp313-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp313-cp313t-macosx_14_0_arm64.whl; tensordict-0.13.0-cp313-cp313t-manylinux_2_28_aarch64.whl; tensordict-0.13.0-cp313-cp313t-manylinux_2_28_x86_64.whl; tensordict-0.13.0-cp313-cp313t-win_amd64.whl; tensordict-0.13.0-cp313-cp313-win_amd64.whl

Tags

More Artificial Intelligence packages