Rate this Page
★ ★ ★ ★ ★
tutorials/tensordict_shapes

Note

Go to the end to download the full example code.

Changing the shape of a TensorDict#

Author: Tom Begley

A TensorDict has the shape operations of a tensor, such as reshape, split, stack, and permute. Each call changes every tensor in it in the same way, so you write the call once, and the tensors stay aligned. This tutorial is for readers who know these operations on a torch.Tensor.

What you will learn:

  • How a shape operation on a tensordict applies to the batch dimensions of every entry.

  • How to split, unbind, stack, and concatenate tensordicts.

  • The difference between a lazy stack and a dense stack.

  • How to change the shape of a tensordict for the length of a with block.

When you create a TensorDict, you give it a batch_size: the leading dimensions that all the entries share. An entry can have more dimensions after them (see Batch size, shapes, and indexing). The shape operations below change the batch dimensions of every entry, and keep the other dimensions.

import torch
from tensordict import TensorDict

torch.manual_seed(0)

<torch._C.Generator object at 0x7ff3e2bf4290>

Indexing#

An index applies to the batch dimensions. Each entry of the TensorDict is indexed in the same way. The Slicing, indexing, and masking tutorial covers indexing in detail.

a = torch.rand(3, 4)
b = torch.rand(3, 4, 5)
tensordict = TensorDict({"a": a, "b": b}, batch_size=[3, 4])

indexed_tensordict = tensordict[:2, 1]
assert indexed_tensordict["a"].shape == torch.Size([2])
assert indexed_tensordict["b"].shape == torch.Size([2, 5])

Reshaping#

TensorDict.reshape works like torch.Tensor.reshape(). It reshapes the batch dimensions of every entry, and keeps the other dimensions, as the shape of b below shows. It also updates the batch_size attribute.

reshaped_tensordict = tensordict.reshape(-1)
assert reshaped_tensordict.batch_size == torch.Size([12])
assert reshaped_tensordict["a"].shape == torch.Size([12])
assert reshaped_tensordict["b"].shape == torch.Size([12, 5])

Splitting#

TensorDict.split works like torch.Tensor.split(). It splits the TensorDict into chunks. Each chunk is a TensorDict with the same keys as the original one. Its entries are views of the entries of the original TensorDict: they share memory with them.

chunks = tensordict.split([3, 1], dim=1)
assert chunks[0].batch_size == torch.Size([3, 3])
assert chunks[1].batch_size == torch.Size([3, 1])
torch.testing.assert_close(chunks[0]["a"], tensordict["a"][:, :-1])

Note

If a function or method takes a dim argument, a negative dim counts from the end of the batch_size of the TensorDict that you call it on. If the tensordict contains nested tensordicts with more batch dimensions, a negative dim still counts from the end of the batch dimensions of the root tensordict.

>>> tensordict = TensorDict(
...     {
...         "a": torch.rand(3, 4),
...         "nested": TensorDict({"b": torch.rand(3, 4, 5)}, [3, 4, 5])
...     },
...     [3, 4],
... )
>>> # dim=-2 is the first dimension for all entries: the root
>>> # TensorDict has 2 batch dimensions, the nested TensorDict has 3
>>> chunks = tensordict.split([2, 1], dim=-2)
>>> assert chunks[0].batch_size == torch.Size([2, 4])
>>> assert chunks[0]["nested"].batch_size == torch.Size([2, 4, 5])

In this example, TensorDict.split gives the same result as with dim=tensordict.batch_dims - 2.

Unbinding#

TensorDict.unbind works like torch.Tensor.unbind(). It removes a dimension and returns a tuple of all the slices along that dimension.

slices = tensordict.unbind(dim=1)
assert len(slices) == 4
assert all(s.batch_size == torch.Size([3]) for s in slices)
torch.testing.assert_close(slices[0]["a"], tensordict["a"][:, 0])

Stacking and concatenating#

torch.cat() and torch.stack() accept lists of TensorDict instances.

Stacking#

You can stack tensordicts in two ways. torch.stack() copies the values into new tensors and returns a TensorDict. lazy_stack() keeps the original tensordicts in a list and returns a LazyStackedTensorDict, which stacks a value only when you read it. Use a lazy stack when the tensordicts have different shapes, devices, or keys, or when you read only a few keys of the stack. The Stacking and variable-size data page explains both kinds of stacks.

from tensordict import lazy_stack

cloned_tensordict = tensordict.clone()
stacked_tensordict = lazy_stack([tensordict, cloned_tensordict], dim=0)
print(stacked_tensordict)

LazyStackedTensorDict(
    fields={
        a: Tensor(shape=torch.Size([2, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([2, 3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
    exclusive_fields={
    },
    batch_size=torch.Size([2, 3, 4]),
    device=None,
    is_shared=False,
    stack_dim=0)

Like torch.stack() on tensors, torch.stack() on tensordicts returns a dense TensorDict:

dense_stack = torch.stack([tensordict, cloned_tensordict])
assert isinstance(dense_stack, TensorDict)

If you index a LazyStackedTensorDict along the stack dimension, you get back the original TensorDict.

assert stacked_tensordict[0] is tensordict
assert stacked_tensordict[1] is cloned_tensordict

If you read a key of the LazyStackedTensorDict, it stacks the values of this key. If the value is a nested TensorDict, you get another LazyStackedTensorDict.

assert stacked_tensordict["a"].shape == torch.Size([2, 3, 4])

Note

A lazy stack stacks a value each time you read it. If you read a value of the stack more than once, convert the LazyStackedTensorDict to a dense TensorDict first, with LazyStackedTensorDict.to_tensordict or LazyStackedTensorDict.contiguous.

>>> assert isinstance(stacked_tensordict.to_tensordict(), TensorDict)
>>> assert isinstance(stacked_tensordict.contiguous(), TensorDict)

The result holds the stacked values, and reading a value does not stack it again.

Concatenating#

Concatenation is never lazy. torch.cat() on a list of TensorDict instances returns a TensorDict. Each of its entries concatenates the entries with the same key in the tensordicts of the list.

concatenated_tensordict = torch.cat([tensordict, cloned_tensordict], dim=0)
assert isinstance(concatenated_tensordict, TensorDict)
assert concatenated_tensordict.batch_size == torch.Size([6, 4])
assert concatenated_tensordict["b"].shape == torch.Size([6, 4, 5])

Expanding#

TensorDict.expand expands all the entries of a TensorDict. As with torch.Tensor.expand(), the expanded entries are views: they share memory with the original entries. Call clone() on the result before you write into it.

exp_tensordict = tensordict.expand(2, *tensordict.batch_size)
assert exp_tensordict.batch_size == torch.Size([2, 3, 4])
torch.testing.assert_close(exp_tensordict["a"][0], exp_tensordict["a"][1])

Squeezing and unsqueezing#

The squeeze() and unsqueeze() methods remove or add a batch dimension of size one.

tensordict = TensorDict({"a": torch.rand(3, 1, 4)}, [3, 1, 4])
squeezed_tensordict = tensordict.squeeze()
assert squeezed_tensordict["a"].shape == torch.Size([3, 4])
print(squeezed_tensordict, end="\n\n")

unsqueezed_tensordict = tensordict.unsqueeze(-1)
assert unsqueezed_tensordict["a"].shape == torch.Size([3, 1, 4, 1])
print(unsqueezed_tensordict)

TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)

TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 1, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3, 1, 4, 1]),
    device=None,
    is_shared=False)

Note

unsqueeze(), squeeze(), view(), permute(), and transpose() apply the operation to the entries when you call them, and return a TensorDict.

These methods apply only to the batch dimensions. The other dimensions of the entries do not change:

tensordict = TensorDict({"a": torch.rand(3, 1, 1, 4)}, [3, 1])
squeezed_tensordict = tensordict.squeeze()
# only one of the two dimensions of size one is removed: the other
# one is not a batch dimension
assert squeezed_tensordict["a"].shape == torch.Size([3, 1, 4])

Viewing#

view() returns a TensorDict of views of the original entries.

tensordict = TensorDict({"a": torch.arange(12)}, [12])
# the entries are viewed at this step
viewed_tensordict = tensordict.view((2, 3, 2))

# "a" is a view of the original tensor
assert viewed_tensordict["a"].shape == torch.Size([2, 3, 2])

Permuting the batch dimensions#

TensorDict.permute reorders the batch dimensions, like torch.permute(). The other dimensions do not change.

Like the other shape operations above, permute applies to the entries when you call it.

tensordict = TensorDict({"a": torch.rand(3, 4), "b": torch.rand(3, 4, 5)}, [3, 4])
# swap the batch dimensions
permuted_tensordict = tensordict.permute([1, 0])

assert permuted_tensordict["a"].shape == torch.Size([4, 3])
assert permuted_tensordict["b"].shape == torch.Size([4, 3, 5])

Reversible operations as context managers#

You can use the result of some reversible operations as a context manager, in a with block. At the end of the block, tensordict applies the inverse operation and writes the changes back into the original tensordict. The shape operations view(), permute(), transpose(), squeeze(), and unsqueeze() work in this way. So do to_module() for functional calls, and lock_() and unlock_(). This example uses transpose:

tensordict = TensorDict({"a": torch.rand(3, 4), "b": torch.rand(3, 4, 5)}, [3, 4])

with tensordict.transpose(1, 0) as tdt:
    tdt.set("c", torch.ones(4, 3))  # the batch size of tdt is [4, 3]

After the block, the original tensordict has the "c" entry, with the batch dimensions in the original order:

assert (tensordict.get("c") == 1).all()
assert tensordict["c"].shape == torch.Size([3, 4])

Gathering values#

TensorDict.gather gathers values along a batch dimension with an index tensor, like torch.gather().

index = torch.randint(4, (3, 4))
gathered_tensordict = tensordict.gather(dim=1, index=index)
print("index:\n", index, end="\n\n")
print("tensordict['a']:\n", tensordict["a"], end="\n\n")
print("gathered_tensordict['a']:\n", gathered_tensordict["a"], end="\n\n")

index:
 tensor([[0, 1, 2, 1],
        [1, 3, 1, 0],
        [0, 0, 3, 3]])

tensordict['a']:
 tensor([[0.0500, 0.4663, 0.9397, 0.2961],
        [0.9515, 0.6811, 0.0488, 0.8163],
        [0.4423, 0.2768, 0.8998, 0.0960]])

gathered_tensordict['a']:
 tensor([[0.0500, 0.4663, 0.9397, 0.4663],
        [0.6811, 0.8163, 0.6811, 0.9515],
        [0.4423, 0.4423, 0.0960, 0.0960]])

Conclusion#

You can now index, reshape, split, unbind, expand, squeeze, unsqueeze, view, permute, and gather a TensorDict with one call each. Each call changes the batch dimensions of every entry and keeps the other dimensions. You can stack tensordicts densely with torch.stack(), or lazily with lazy_stack() when their shapes differ. You can also change the shape for the length of a with block, and get the changes back in the original tensordict.

The same operations work on a TensorClass, the class form of a tensordict with typed fields (see TensorClass).

Further reading#

Total running time of the script: (0 minutes 0.009 seconds)

Gallery generated by Sphinx-Gallery

On this page
Show Source
PyTorch Libraries