reference/generated/tensordict.TensorDictBase.expand

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.expand

TensorDictBase.expand(*args, **kwargs) → Self

Expands each tensor of the tensordict according to torch.Tensor.expand(), ignoring the feature dimensions.

The new shape is the new batch size. As for torch.Tensor.expand(), it can add leading dimensions and expand the dimensions of size 1. The other dimensions must keep their size. The result holds views of the original tensors: no memory is allocated, and the expanded leaves are not contiguous. Call clone() or contiguous() before you write into them.

Supports iterables to specify the shape.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td_expand = td.expand(10, 3, 4)
>>> assert td_expand.shape == torch.Size([10, 3, 4])
>>> assert td_expand.get("a").shape == torch.Size([10, 3, 4, 5])

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

Tutorials that use expand

Changing the shape of a TensorDict

Changing the shape of a TensorDict