TensorDict

A tensorclass is a dataclass-like container that inherits all of TensorDict's capabilities (indexing, reshaping, to(device), stack/cat, memory-mapped serialization, and torch.compile support) while providing typed attribute access instead of string keys. This allows for IDE/type-checker support and constrained field sets.

There are two ways to define a tensorclass:

1. Inheritance-based (Recommended): Inherit from TensorClass. This is best for static type-checkers and explicit configuration.
2. Decorator-based: Use the @tensorclass decorator. This is useful for migrating existing @dataclass code.

When a batch_size is provided, tensor fields are indexed elementwise, while non-tensor fields are preserved.

from __future__ import annotations
from typing import Optional
import torch
from tensordict import TensorClass

class MyData(TensorClass):
    floatdata: torch.Tensor
    intdata: torch.Tensor
    non_tensordata: str
    nested: Optional[MyData] = None

    def check_nested(self):
        assert self.nested is not None

data = MyData(
    floatdata=torch.randn(3, 4, 5),
    intdata=torch.randint(10, (3, 4, 1)),
    non_tensordata="test",
    batch_size=[3, 4],
)

# Indexing works on tensor fields and nested tensorclasses
indexed = data[:2]

When you use memmap_ or memmap with a prefix, TensorDict creates a directory tree that mirrors your data structure. Each tensor is stored as a .memmap file, and metadata required to reconstruct the object (like device, batch_size, and subtypes) is stored in meta.json files.

This allows TensorDict.load_memmap to reconstruct complex nested structures, including tensorclass objects and nested TensorDicts, even if they have different types than their parents.

from tensordict import TensorDict, tensorclass, TensorDictBase
from tensordict.utils import print_directory_tree
import torch
import tempfile

@tensorclass
class MyClass:
    data: torch.Tensor
    metadata: str

# Setup complex nested data
td_list = [TensorDict({"item": i}, batch_size=[]) for i in range(4)]
tc = MyClass(torch.randn(3), metadata="some text", batch_size=[])
data = TensorDict({"td_list": torch.stack(td_list), "tensorclass": tc}, [])

with tempfile.TemporaryDirectory() as tempdir:
    data.memmap_(tempdir)
    loaded_data = TensorDictBase.load_memmap(tempdir)
    assert (loaded_data == data).all()
    print_directory_tree(tempdir)

Choose between TypedTensorDict and TensorClass based on your data requirements:

Summary:

• Use TypedTensorDict for typed pipelines, progressive state accumulation, and wrapping persistent backends.
• Use TensorClass for non-tensor metadata, custom __init__ logic, or if your codebase heavily uses the @tensorclass decorator.

TensorDict is compatible with torch.func transforms like jacrev, jacfwd, and hessian. When applied to a function that accepts and returns a TensorDict, these transforms produce structured Jacobians/Hessians.

Each entry J["output_key", "input_key"] represents the Jacobian block for that specific pair of keys.

Important Constraint: All tensors in the TensorDict must have at least one non-batch (feature) dimension (tensor.ndim > len(batch_size)). If a tensor's shape equals the batch_size, the Jacobian basis vectors become ambiguous with the batch dimensions, causing errors in torch.func transforms.

To fix this, either:

1. Set batch_size=[] (treating all dimensions as features).
2. Add a trailing dimension using .unsqueeze(-1).

import torch
from torch.func import jacrev
from tensordict import TensorDict

td = TensorDict({"a": torch.randn(3, 2), "b": torch.randn(3, 4)}, batch_size=[3])

def f(td):
    return TensorDict(
        {"x": td["a"] ** 2, "y": td["b"] ** 3},
        batch_size=td.batch_size
    )

J = jacrev(f)(td)
print(J.batch_size)  # torch.Size([3])
print(J["x", "a"].shape)  # output_shape + input_shape = (3, 2, 3, 2)

A TensorDict is defined by its batch_size (or shape) and its key-value pairs. Key metadata includes:

• batch_size: The dimensions representing the batch.
• device: The device where tensors are stored.
• is_memmap / is_shared: Shared memory status.
• names: Dimension names.
• lock: The modification lock status.

When initializing, the batch_size must be compliant with the first dimensions of each tensor. You can also specify a device during initialization; all subsequent write operations will automatically cast tensors to that device.

import torch
from tensordict import TensorDict

# Defining a TensorDict with a specific batch_size and device
data = TensorDict({
    "key 1": torch.ones(3, 4, 5),
    "key 2": torch.zeros(3, 4, 5, dtype=torch.bool),
}, batch_size=[3, 4], device="cuda:0")

# Writing a CPU tensor will automatically cast it to the TensorDict's device (cuda:0)
data["key 3"] = torch.randn(3, 4, device="cpu")
assert data["key 3"].device is torch.device("cuda:0")

Equality Operators

Tensorclasses support == and != operators, including for nested instances. However, these operators do not validate non-tensor/meta data. When comparing tensorclasses, the result is a new tensorclass where:

• Tensor fields contain boolean leaves.
• Non-tensor fields are set to None.

Item Assignment

For performance reasons, item assignment performs an identity check on non-tensor/meta data rather than an equality check. If you assign a slice or item where the non-tensor data differs from the target, a UserWarning is emitted. Users are responsible for ensuring non-tensor data remains in sync.

Concatenation and Stacking

torch.cat and torch.stack work on tensorclasses but do not validate non-tensor/meta fields. The operation runs on the tensor leaves, and the non-tensor data from the first instance in the list is preserved. If inputs disagree on a non-tensor field, the output silently follows the first one.

# Equality behavior example
print(data == data2)
# Output contains boolean tensors for data and None for non-tensor fields

# torch.cat behavior: the first instance's meta data wins
data2.non_tensordata = "test_new"
stack_tc = torch.cat([data, data2], dim=0)
assert stack_tc.non_tensordata == "test"  # data's value wins

If you need to store tensors in a TensorDict that should not be affected by batch-size-related shape operations (like reshape, unbind, or split), wrap them in an UnbatchedTensor.

This is useful for configuration tensors, masks, or parameters shared across an entire batch.

Contract of UnbatchedTensor

• Tensor-like: Behaves like a torch.Tensor for arithmetic, device transfers (to), and gradient computation.
• Shape-operation pass-through: Operations like reshape, view, unbind, split, squeeze, etc., return copies of the underlying tensor without modifying it, even when called on the parent TensorDict.
• Scalar conversion: Follows standard PyTorch behavior for float(), int(), etc.

Limitations

• Memory-mapped serialization is not supported for UnbatchedTensor.
• When stacking TensorDicts containing UnbatchedTensor entries with different underlying data, only the first element's data is kept (a warning is emitted).

from tensordict import TensorDict, UnbatchedTensor
import torch

td = TensorDict(
    a=torch.randn(2, 3),
    config=UnbatchedTensor(torch.tensor([1.0, 2.0, 3.0])),
    batch_size=[2, 3],
)

# Shape operations on the TensorDict leave the UnbatchedTensor storage unchanged
reshaped = td.reshape(6)
assert reshaped["config"].data_ptr() == td["config"].data_ptr()

parts = td.unbind(0)
assert parts[0]["config"].data_ptr() == parts[1]["config"].data_ptr()

# Pointwise arithmetic is applied to the underlying data
td2 = td * 2
assert torch.equal(td2["config"], torch.tensor([2.0, 4.0, 6.0]))

When storing heterogeneous data (strings, Python objects, etc.) in a TensorDict, you should use the NonTensorData wrapper. Each backend handles the serialization of this data differently:

• memmap: Serializes via NonTensorData into meta.json (JSON) or other.pickle (pickle fallback).
• HDF5: Uses HDF5 string or opaque datasets; NonTensorData is applied automatically on read.
• Zarr: Uses a JSON or pickle payload within a marked uint8 array; NonTensorData is applied automatically on read.
• Redis: Transparently serializes as JSON (falling back to pickle for non-JSON-serializable objects) via metadata hashes.

To ensure compatibility across these backends, wrap your non-tensor values in NonTensorData during assignment.

from tensordict import TensorDict, NonTensorData
import torch

td = TensorDict(
    obs=torch.randn(4, 3),
    label=NonTensorData(data="cat", batch_size=[4]),
    batch_size=[4],
)

# Example with memmap
td_mm = td.memmap_("/tmp/example")
loaded = TensorDict.load_memmap("/tmp/example")
print(loaded["label"].data)  # Output: 'cat'