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 a single tensordict along the given dimension.

This call is equivalent to calling torch.stack() but is compatible with torch.compile. The tensordicts must have the same batch size and the same keys. Otherwise, the call raises RuntimeError. 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.

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.