reference/generated/tensordict.TensorDictBase.swapaxes

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.swapaxes

TensorDictBase.swapaxes(axis0: int, axis1: int)

Interchange two axes of the tensordict.

This is an alias for transpose().

Parameters:

• axis0 (int) – First axis.

• axis1 (int) – Second axis.

Returns:

a new tensordict with the axes swapped.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.randn(3, 4, 5)}, batch_size=[3, 4])
>>> print(td.swapaxes(0, 1).shape)
torch.Size([4, 3])

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