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 into a new entry and possibly remove the original values.

Parameters:

keys (sequence of NestedKey) – entries to stack.

Keyword Arguments:

• out_key (NestedKey) – new key name for the stacked inputs.

• keep_entries (bool, optional) – if False, entries in keys will be deleted. Defaults to False. Entries are deleted only after the result has been written. If stacking fails, or if the result cannot be written, every input entry is kept with its original value.

• dim (int, optional) – the dimension along which the stack must occur. Defaults to 0.

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.