reference/generated/tensordict.TensorDictBase.stack

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.stack

classmethod TensorDictBase.stack(input, dim: int = 0, *, out=None)

Stacks tensordicts into one dense tensordict along a new dimension dim.

It is equivalent to torch.stack(), and also works under torch.compile(). The tensordicts must have the same batch size and the same keys, or a RuntimeError is raised. To stack tensordicts with different keys or shapes, use lazy_stack(). See Stacking and variable-size data.

Parameters:

• input (sequence of TensorDictBase) – the tensordicts to stack.

• dim (int, optional) – the new batch dimension. Defaults to 0.

Keyword Arguments:

out (TensorDictBase, optional) – a tensordict to write the result into. Defaults to None.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> a = TensorDict(x=torch.zeros(3), batch_size=[3])
>>> b = TensorDict(x=torch.ones(3), batch_size=[3])
>>> TensorDict.stack([a, b]).batch_size
torch.Size([2, 3])

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