# TensorDict: a dict of tensors that behaves like one tensor, and the batch size contract behind it

> TensorDict is the PyTorch container that holds a group of tensors sharing their leading dimensions, so slicing, moving to a device and compiling apply to all of them at once, with a TensorClass variant for typed fields and methods. This looks at what it does well and what it refuses to blur.

**pytorch/tensordict** — TensorDict is a pytorch dedicated tensor container.

- Repository: https://github.com/pytorch/tensordict
- Stars: 1,042 · Forks: 146
- Language: Python
- License: MIT
- Published: 2026-10-19 · Updated: 2026-10-19 · Language: en
- Canonical page: https://hysenlabs.com/projects/pytorch-tensordict

## The problem is repeating yourself over every tensor

TensorDict is described as a pytorch dedicated tensor container, and the readme puts it as a dict of strings to torch.Tensor that behaves like one tensor.

The problem it names is specific. PyTorch code that handles groups of tensors, such as a training batch, a particle system, or a batch of graphs, usually repeats every operation for every tensor, and nested dicts need recursion. So a training loop ends up looking like a pile of parallel lists, where moving the batch to the GPU means writing the same call once per field, and slicing it means doing it once per field, and it is easy to forget one.

The TensorDict holds the group as one object instead:

```python
batch = TensorDict(
    {"image": images, "label": labels, "text": {"tokens": tokens, "mask": mask}},
    batch_size=[32],
)

batch = batch.to("cuda")                      # moves every tensor, nested ones too
train, val = batch[:24], batch[24:]           # slices every tensor
shuffled = batch[torch.randperm(len(batch))]  # one permutation for all of them
pairs = torch.stack([batch, batch])           # batch_size [2, 32]
batch["text", "mask"]                         # nested keys are tuples
flat = batch.flatten_keys(".")                # keys: image, label, text.tokens, text.mask
```

The mechanism is that the tensors share their leading dimensions, the batch size, so indexing, reshaping, stacking, `.to(device)`, saving and `torch.compile` apply to all of them at once. That is the whole design, and everything else in the library is a consequence of it.

The audience is PyTorch developers working with structured batches, and the readme names where it is used: physics through PhysicsNeMo, robotics through MuJoCo Playground and ProtoMotions, reinforcement learning through TorchRL, LLM post-training through verl and ROLL, and genomics through Medaka, with more than 2,000 other repositories depending on it.

## The batch size is a contract, and it fails loudly

The batch size is not a hint. Every tensor added to a TensorDict must have the batch dimensions, and the readme shows what happens when it does not.

```python
batch["scores"] = torch.ones(31)
# RuntimeError: batch dimension mismatch, got self.batch_size=torch.Size([32]) and
# value.shape=torch.Size([31]) for key 'scores'. The leading dimensions of a value
# must match the batch size.
```

That error message is unusually informative. It names the key, it prints both shapes, and it states the rule. Which is what you want, because the alternative failure mode in PyTorch code, where a list of tensors silently drifts out of sync with the batch, produces a wrong answer rather than an exception.

The mechanics of a nested structure are worth spelling out, because they are where a dict stops being a dict. Nested keys are tuples, so `batch["text", "mask"]` reaches the nested mask, and `flatten_keys(".")` turns the nested structure into a flat one with dot separated keys, which is what you want when a downstream consumer expects flat named entries.

```python
import torch
from tensordict import TensorDict

batch = TensorDict(
    {"images": torch.randn(32, 3, 64, 64), "labels": torch.randint(10, (32,))},
    batch_size=[32],  # the leading dimension that all tensors share
)

print(batch[:8].batch_size)                 # torch.Size([8])
print(batch.reshape(4, 8)["images"].shape)  # torch.Size([4, 8, 3, 64, 64])
```

Note the last line. A reshape of the container reshapes every tensor, and indexing the result gives the reshaped image tensor. That is the property that makes the container worth its overhead in a training loop.

## TensorClass, where a method runs on one item or on a whole batch

TensorDict has keys at run time, like a dict. TensorClass is the other mode, declared fields and methods, like a dataclass.

```python
import torch
from tensordict import TensorClass

class Particles(TensorClass):
    pos: torch.Tensor   # [*batch, 3]
    vel: torch.Tensor   # [*batch, 3]
    mass: torch.Tensor  # [*batch]

    def step(self, force: torch.Tensor, dt: float) -> "Particles":
        vel = self.vel + dt * force / self.mass.unsqueeze(-1)
        return Particles(pos=self.pos + dt * vel, vel=vel, mass=self.mass, batch_size=self.batch_size)
```

The claim the documentation makes is that a method written once runs on one item or on a batch of any shape, and compiles. The example shows both halves. The batch is 64 systems of 1,000 particles each, and `systems.step(...)` updates all 64, while `systems[0].step(...)` updates one system and returns a `Particles` with a batch size of `torch.Size([1000])`.

That last detail is the one worth checking in your own code. Slicing the container narrows the batch dimension, and a method that returns a new instance has to pass `batch_size` along explicitly, which is why the return statement carries `batch_size=self.batch_size`. The type system does not infer that for you.

The compile path is then a decorator away:

```python
step = torch.compile(Particles.step, fullgraph=True)
systems = step(systems, torch.randn(64, 1000, 3), 0.01)
```

`fullgraph=True` is the strict form, which fails the compile rather than falling back to eager execution with a graph break. For a library whose selling point is that the same method works on a batch and on an item, that is the correct choice, and it is the one that would catch a method that quietly breaks the batching rules.

## Ragged data and the memmap path

Two features handle the cases where the neat column model does not fit, and both are in the readme's collapsible examples.

The first is variable-length data. Sequences of different lengths cannot share a batch dimension, so the container marks the varying dimension with -1:

```python
import torch
from tensordict import TensorDict, lazy_stack

episodes = [TensorDict({"obs": torch.randn(n, 4)}, batch_size=[n]) for n in (3, 5, 2)]
batch = lazy_stack(episodes)      # keeps each sample as it is
print(batch.batch_size)           # torch.Size([3, -1]): -1 marks the dimension that varies
print(batch[1]["obs"].shape)      # torch.Size([5, 4])
print(batch[:, :2]["obs"].shape)  # torch.Size([3, 2, 4]): the first 2 steps of each
```

The name lazy_stack is doing the work: the samples are stacked without padding, and the -1 marks the dimension that varies, so indexing still works, `batch[1]` is the five step episode and `batch[:, :2]` is the first two steps of each.

The second is storage. A dataset can be written to disk and read back lazily:

```python
data.memmap(path)                      # one file per tensor, plus JSON metadata
loaded = TensorDict.load_memmap(path)  # maps the files; reads no data yet
print(loaded[1234]["labels"])          # MemoryMappedTensor(1234)
```

The layout is one file per tensor plus a JSON metadata file, and loading maps the files without reading the data, so indexing a single sample does not pull the whole dataset into memory. For a dataset of images, that is the difference between a data loader that streams and one that does not.

## tensordict.nn, and running an ensemble with vmap

There is a module layer, `tensordict.nn`, for models that read and write named entries rather than positional tensors.

```python
import torch
from torch import nn
from tensordict import TensorDict
from tensordict.nn import TensorDictModule, TensorDictSequential

policy = TensorDictSequential(
    TensorDictModule(nn.Linear(4, 32), in_keys=["obs"], out_keys=["hidden"]),
    TensorDictModule(nn.Linear(32, 2), in_keys=["hidden"], out_keys=["action"]),
)
batch = policy(TensorDict({"obs": torch.randn(10, 4)}, batch_size=[10]))
print(batch["action"].shape)  # torch.Size([10, 2])
```

The `in_keys` and `out_keys` arguments are the mechanism. A module declares which named entries it consumes and which it produces, and the container wires the outputs to the inputs of the next module. So the model reads `obs` and writes `hidden`, the next reads `hidden` and writes `action`, and the result is a TensorDict with the right entries, rather than a tuple of tensors you unpack by hand.

The same idea extends to ensembles, and this is the example that makes the case for the container:

```python
models = [nn.Linear(4, 2) for _ in range(5)]
params = TensorDict.from_modules(*models)  # the parameters of the 5 models, batch_size [5]
base = models[0]

def call(p, x):
    with p.to_module(base):
        return base(x)

out = torch.vmap(call, (0, None))
```

`TensorDict.from_modules` collects the parameters of five models into one container with a batch size of 5, and `to_module` puts a given set of parameters back into a module's structure, so `torch.vmap` runs the same call across the ensemble. A plain dict of parameter dicts would not give you that, because the parameters have to be addressable as tensors.

## The repository behind the library

The tree shows a project with real engineering discipline, which matters because this is a library that sits underneath research code.

At the root there is `ARCHITECTURE.md`, `GETTING_STARTED.md`, `AGENTS.md`, `CONTRIBUTING.md`, a `CODE_OF_CONDUCT.md`, `mypy.ini` and `pytest.ini`, plus `benchmarks/`, `docs/`, `tutorials/`, `scripts/`, `test/` and the `tensordict/` package itself. The build is setuptools with setuptools_scm, and there is a `version.txt`.

The `setup.py` is worth reading for one specific reason: it computes the version from the branch and the short commit SHA, producing a local version of the form `base_version+g<sha>`, and it recognises branches named `release/v<id>`. That is how a 0.x library pins releases: the release branch determines the base version and a local build gets the SHA appended. The `pyproject.toml` also carries the Ruff configuration, with a long ignore list that documents exactly which checks the previous flake8 setup ran and which Ruff rules were deliberately left out, and per-file ignores for `__init__.py` and the test files.

The licence is MIT, and the setup script header records Meta Platforms, Inc. and affiliates as the copyright holder. The release cadence is fast: v0.14.2 on 2026-09-09, v0.14.3 on 2026-10-08, and v0.15.0 on 2026-10-15, with the last push to the repository on the same day.

## Where tensordict is the wrong tool

Four limits are clear.

The first is the batch size contract, which is a runtime error rather than a static one. Every value you put in must share the leading dimensions of the batch. That is the right rule for a batch, and it is the wrong rule for data whose shapes legitimately differ, and the answer to that is `lazy_stack` with a -1 dimension, which is a different code path with different behaviour from ordinary indexing. Data that is neither uniform nor stackable is where this container stops helping.

The second is that it is a container, not a data pipeline. It does not read files, shuffle, or feed a model. It holds a batch that something else produced, and it makes operations on that batch uniform. If the problem is that your dataset loader is slow, a different library is the answer.

The third is the version line. The project is on 0.15.0, and a 0.x library in the PyTorch ecosystem can and does change its API between minors. The release branch mechanism in `setup.py` is the mitigation, pinning a base version to a branch, but a team that upgrades a minor version should expect to read the changelog.

The fourth is the requirements. The readme states pure Python, for Python 3.10+ and PyTorch 2.13+, with nightly builds and extras available. A PyTorch 2.13 requirement is a recent floor, so on an older pinned environment the choice is an older tensordict release or a different container.

None of that makes it a library to avoid. The evidence that it works is in the readme's list of users, PhysicsNeMo, MuJoCo Playground, ProtoMotions, TorchRL, verl, ROLL and Medaka, and the physics and reinforcement learning cases are exactly the ones where a step function has to run over thousands of systems at once.

## Conclusion

TensorDict fits a PyTorch team that keeps repeating the same operation across a group of tensors, especially in reinforcement learning, robotics and physics where a step function has to run over a whole batch of systems. It does not fit a project whose data is naturally ragged without the extra work of the variable-length path, and it is not a replacement for a Dataset or a data loader. Before you adopt it, check the batch size contract against your data, because a tensor whose leading dimension does not match raises at run time rather than at import, and note that the project is still on a 0.x version line where the release branches in `setup.py` are how versions are pinned.

## FAQ

### How do I install tensordict?

Run `pip install tensordict`. The library is pure Python and requires Python 3.10 or later and PyTorch 2.13 or later, with nightly builds and extras documented on the installation page of the documentation site.

### What is the batch size contract in a TensorDict?

Every tensor added must have the leading dimensions of the shared batch size, otherwise assignment raises a RuntimeError naming the key and both shapes. The batch size is declared when the container is created, as `batch_size=[32]` for example.

### What is the difference between TensorDict and TensorClass?

TensorDict has keys at run time, like a dict, while TensorClass declares typed fields and methods, like a dataclass. A TensorClass method written once runs on a single item or on a batch of any shape, and can be compiled with torch.compile.

### How do I handle sequences of different lengths in tensordict?

Use lazy_stack, which stacks samples without padding and marks the varying dimension with -1 in the batch size, so `batch[1]` returns the episode of its own length and `batch[:, :2]` returns the first two steps of each.

### Can I save a tensordict to disk and read samples lazily?

Yes. `data.memmap(path)` writes one file per tensor plus JSON metadata, and `TensorDict.load_memmap(path)` maps those files without reading any data, so indexing a single sample returns a MemoryMappedTensor rather than loading the whole dataset.

## Sources

- [Issues](https://github.com/pytorch/tensordict/issues)
- [License: MIT](https://github.com/pytorch/tensordict/blob/main/LICENSE)
- [pytorch/tensordict on GitHub](https://github.com/pytorch/tensordict)
- [README](https://github.com/pytorch/tensordict/blob/main/README.md)
- [Releases](https://github.com/pytorch/tensordict/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/pytorch-tensordict