reference/generated/tensordict.TensorDictBase.stack_tensors

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.stack_tensors

TensorDictBase.stack_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → Self

Stacks entries of the tensordict into a new entry, and removes them by default.

Parameters:

*keys (NestedKey) – the entries to stack, with torch.stack().

Keyword Arguments:

• out_key (NestedKey) – the key of the result.

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

• keep_entries (bool, optional) – if False, removes the entries of keys once the result is written. If the stack or the write fails, every entry is kept with its original value. Defaults to False.

Returns:

self.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(()), b=torch.ones(()))
>>> _ = td.stack_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

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