reference/generated/tensordict.utils.expand_right

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

expand_right

tensordict.utils.expand_right(tensor: Tensor, shape: Sequence[int]) → Tensor

Expands a tensor on the right to shape.

Broadcasting adds dimensions on the left; this function adds them on the right instead, then expands every dimension of size one. The leading dimensions of tensor must match shape or be of size one. The result is a view of tensor, as with torch.Tensor.expand().

Parameters:

• tensor (Tensor) – the tensor to expand.

• shape (Size or sequence of int) – the shape of the result.

Returns:

A view of tensor with shape shape.

Examples

>>> import torch
>>> from tensordict.utils import expand_right
>>> x = torch.arange(3)
>>> expand_right(x, (3, 2))
tensor([[0, 0],
        [1, 1],
        [2, 2]])
>>> expand_right(torch.zeros(3, 4), (3, 4, 5)).shape
torch.Size([3, 4, 5])