reference/generated/tensordict.TensorDictBase.stack_from_tensordict

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.stack_from_tensordict

TensorDictBase.stack_from_tensordict(dim: int = 0, *, sorted: bool | List[NestedKey] | None = None, out: Tensor | None = None) → Tensor

Stacks all the leaves of the tensordict into one tensor, with torch.stack().

The leaves must have the same shape.

Parameters:

dim (int, optional) – the new dimension, counted in the dimensions of the result. Defaults to 0.

Keyword Arguments:

• sorted (bool or list of NestedKey, optional) – the order of the leaves. True sorts the keys alphabetically, and False or None keeps the order of insertion. A list of keys gives the order, and must hold every leaf key, or a RuntimeError is raised. Defaults to None.

• out (Tensor, optional) – a tensor to write the result into. Defaults to None.

Returns:

The stacked tensor.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(2), b=torch.ones(2), batch_size=[2])
>>> td.stack_from_tensordict(dim=1)
tensor([[0., 1.],
        [0., 1.]])

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