TensorDictBase.roll#

TensorDictBase.roll(shifts: int | tuple[int, ...], dims: int | tuple[int, ...] = None, *, inplace: bool = False)#

Returns a tensordict with the elements rolled along batch dimensions, as torch.roll() does.

Elements shifted past the last position come back at the first one. The batch size and the names do not change, and the feature dimensions are not rolled. The data is copied.

It can be used as a context manager, as unsqueeze() can: when the with block exits, the entries of the block are written back to the tensordict with the inverse operation.

Parameters:

shifts (int or tuple of int) – the number of places to shift the elements by. With a tuple, dims is a tuple of the same size, and each dimension is rolled by its own shift.
dims (int or tuple of int, optional) – the batch dimensions to roll. Defaults to None, which rolls the flattened batch and restores its shape.

Keyword Arguments:

inplace (bool, optional) – if True, the tensordict keeps its identity and keys, its leaves are replaced one by one, and the call returns self. See In-place operations. Defaults to False.

Returns:

The rolled tensordict, or self with inplace=True.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.arange(6).view(2, 3)}, batch_size=[2, 3])
>>> print(td["a"])
tensor([[0, 1, 2],
        [3, 4, 5]])
>>> print(td.roll(1, 0)["a"])
tensor([[3, 4, 5],
        [0, 1, 2]])

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

TensorDictBase.roll#

Docs

Tutorials

Resources