reference/generated/tensordict.TensorDictBase.cat_from_tensordict

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.cat_from_tensordict

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

Concatenates all the leaves of the tensordict into one tensor, with torch.cat().

Parameters:

dim (int, optional) – the dimension to concatenate along, counted in the dimensions of the leaves. 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 concatenated tensor.

See also

pack(), which flattens leaves of any shape into one tensor and returns the layout to rebuild the tensordict.

Examples

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

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