Filling a TensorDict row by row#

A simulation loop often stores results at each step, such as the positions and the velocities of its particles. To collect them, you usually allocate one tensor per quantity in advance, then write each step into it. A TensorDict does the allocation for you: the first write with an index creates every tensor at its full size. This tutorial is for readers who fill a buffer one step at a time.

What you will learn:

How to create an empty tensordict with a batch size and no keys.
What tensordict allocates when you write the first row with an index.
How the later rows are written into the same tensors.

Here is a function that returns a TensorDict with a batch size of [3]:

import torch
from tensordict import TensorDict

torch.manual_seed(0)


def make_tensordict():
    return TensorDict({"a": torch.rand(3), "b": torch.rand(3, 4)}, [3])

Call this function N times, and write the results into one TensorDict with batch_size=[N, 3]:

N = 10
tensordict = TensorDict({}, batch_size=[N, 3])

for i in range(N):
    tensordict[i] = make_tensordict()

print(tensordict)

TensorDict(batch_size=[10,3], fields={a: f32[], b: f32[4]})

tensordict has batch_size=[N, 3] and no keys. When the first iteration writes tensordict[0], tensordict creates each missing key as a zero-filled tensor. Its shape is the batch size [N, 3] followed by the other dimensions of the value from make_tensordict. Here "a" gets shape torch.Size([10, 3]) and "b" gets shape torch.Size([10, 3, 4]). The later iterations write into these tensors in place. Rows that the loop does not write stay zero.

The steps below go through the loop one iteration at a time. First, create an empty TensorDict:

N = 10
tensordict = TensorDict({}, batch_size=[N, 3])
print(tensordict)

TensorDict(batch_size=[10,3], fields={})

After the first iteration, tensordict has tensors for "a" and "b". These tensors are zero, except for the first row, which holds the random values that you wrote.

random_tensordict = make_tensordict()
tensordict[0] = random_tensordict

assert (tensordict[1:] == 0).all()
assert (tensordict[0] == random_tensordict).all()

print(tensordict)

TensorDict(batch_size=[10,3], fields={a: f32[], b: f32[4]})

The next iterations write into the same tensors, in place.

a = tensordict["a"]
random_tensordict = make_tensordict()
tensordict[1] = random_tensordict

# the same tensor is stored under "a", but its values changed
assert tensordict["a"] is a
assert (tensordict[:2] != 0).all()

Conclusion#

You can now fill a TensorDict one row at a time, without allocating its tensors yourself. Create it with a batch size and no keys. The first write allocates a zero-filled tensor for each key, at the full batch size. The later writes go into the same tensors.

To allocate all the keys in advance, from their shapes and dtypes, without a first row, use from_schema(). It can also allocate the tensors on disk.

Filling a TensorDict row by row#

Conclusion#

Further reading#

Docs

Tutorials

Resources