reference/generated/tensordict.TensorDictBase.permute

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.permute

TensorDictBase.permute(*args, **kwargs)

Returns a view of a tensordict with the batch dimensions permuted according to dims.

Parameters:

• *dims (int) – the new ordering of the batch dims of the tensordict. Alternatively, a single iterable of integers can be provided.

• dims (list of int) – alternative way of calling permute(…).

Returns:

a new tensordict with the batch dimensions in the desired order.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4])
>>> print(tensordict.permute([1, 0]))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 3, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)
>>> print(tensordict.permute(1, 0))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 3, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)
>>> print(tensordict.permute(dims=[1, 0]))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 3, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)

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

Tutorials that use permute

Changing the shape of a TensorDict

Changing the shape of a TensorDict