reference/generated/tensordict.TensorDictBase.broadcast_to

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.broadcast_to

TensorDictBase.broadcast_to(shape: tuple[int, ...])

Returns the tensordict broadcast to the batch size shape, as torch.broadcast_to() does.

As for expand(), the new shape can add leading dimensions and expand the dimensions of size 1, and -1 keeps the size of a dimension. The leaves of the result are views of the original leaves, and the result has no names.

Parameters:

shape (tuple of int or Size) – the new batch size.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.arange(3)}, batch_size=[3])
>>> print(td.broadcast_to((2, 3)).shape)
torch.Size([2, 3])

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