reference/generated/tensordict.TensorDictBase.flip

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.flip

TensorDictBase.flip(dims: int | tuple[int, ...])

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

The batch size and the names do not change, and the feature dimensions are not flipped. 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:

dims (int or tuple of int) – the batch dimensions to flip.

Raises:

IndexError – if a dimension is not a batch dimension.

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.flip(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.