---
id: tensordict
version: "0.13.0"
license: BSD
license_treatment: permissive
maintenance: active
---

# tensordict — TensorDict is a pytorch dedicated tensor container.

License: permissive · Maintenance: active · Downloads: 2.6M/mo

## 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 above — 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

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

- License: BSD (permissive)
- Python support: supports_current
- Install friction: medium
- Maintenance: active
- Downloads: 2.6M/month (top 5,000 on PyPI)
- Known vulnerabilities: none known

## Tags

batched nested tensor container, dict-like tensor operations, structured tensor batching, pytorch tensor dictionary, synchronized nested tensor slicing, batch-aware data structure, tensor collection framework, pytorch-native, batching-framework, rl-infrastructure

[View on SkillFed](https://skillfed.io/packages/tensordict) · [View on PyPI](https://pypi.org/project/tensordict/)