reference/generated/tensordict.TensorDictBase.maybe_dense_stack

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.maybe_dense_stack

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

Attempts to make a dense stack of tensordicts, and falls back on lazy stack when required.

The result is a TensorDict if the tensordicts have the same keys and shapes, and a LazyStackedTensorDict otherwise. See Stacking and variable-size data, and tensordict.maybe_dense_stack() for the arguments.

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])
>>> c = TensorDict(y=torch.ones(3), batch_size=[3])
>>> type(TensorDict.maybe_dense_stack([a, b])).__name__
'TensorDict'
>>> type(TensorDict.maybe_dense_stack([a, c])).__name__
'LazyStackedTensorDict'

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