reference/generated/tensordict.TensorDictBase.repeat_interleave

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.repeat_interleave

TensorDictBase.repeat_interleave(repeats: Tensor | int, dim: int | None = None, *, output_size: int | None = None, inplace: bool = False) → Self

Repeats each element of the batch along dim, as torch.Tensor.repeat_interleave() does.

Each element is repeated next to itself, as with numpy.repeat(). To repeat the whole batch instead, as with numpy.tile(), use repeat(). The data is copied. With a dim, the names are kept.

Parameters:

• repeats (Tensor or int) – the number of repetitions of each element, broadcast to the size of dim.

• dim (int, optional) – the batch dimension to repeat along. Defaults to None, which flattens the batch dimensions first, so the result has one batch dimension. A dim past the batch dimensions warns, and is used on the dimensions of the leaves. Starting with TensorDict 0.17, the default will change to raising an IndexError, as for a tensor.

Keyword Arguments:

• output_size (int, optional) – the size of the result along dim (the sum of repeats). Passing it avoids the device synchronization that computes it. Defaults to None, which computes it.

• inplace (bool, optional) – if True, the tensordict keeps its identity and keys, its leaves are replaced one by one, and the call returns self. It needs a dim and at least one batch dimension. A lazy stack, or a tensordict that holds one, raises a NotImplementedError: call to_tensordict() first. See In-place operations. Defaults to False.

Returns:

A tensordict with the same batch size, except along dim, or self with inplace=True.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.tensor([[0, 1], [2, 3]]), batch_size=[2, 2])
>>> td.repeat_interleave(2, dim=0)["a"]
tensor([[0, 1],
        [0, 1],
        [2, 3],
        [2, 3]])
>>> td.repeat_interleave(torch.tensor([1, 2]), dim=1)["a"]
tensor([[0, 1, 1],
        [2, 3, 3]])
>>> td.repeat_interleave(2).batch_size
torch.Size([8])

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