reference/generated/tensordict.TensorDictBase.add

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.add

TensorDictBase.add(other: TensorCollection | Tensor, *, alpha: float | None = None, default: str | Tensor | TensorCollection | None = None) → Self

Returns a new tensordict with other, scaled by alpha, added to every leaf.

td + other and other + td call this method.

\[\text{out}_i = \text{input}_i + \text{alpha} \times \text{other}_i\]

Parameters:

other (TensorDictBase, Tensor or Number) – the value to add. A tensordict is paired with the tensordict key by key. A tensor is aligned with the batch dimensions of each leaf and expanded over its feature dimensions. Batch sizes that differ are broadcast together.

Keyword Arguments:

• alpha (Number, optional) – the multiplier for other. Defaults to None, which adds other as it is.

• default (Tensor or str, optional) – what to do with the keys that only one of the two tensordicts has. None requires the same keys and raises a KeyError otherwise. "intersection" computes the keys that both have and drops the others. A tensor stands in for the missing entry on either side, and the result has the keys of both. Defaults to None.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(2, 3), b=torch.ones(2), batch_size=[2])
>>> td.add(1, alpha=2)["b"]
tensor([3., 3.])
>>> td.add(torch.tensor([1.0, 2.0]))["a"]  # one value per batch element
tensor([[1., 1., 1.],
        [2., 2., 2.]])
>>> other = TensorDict(b=torch.ones(2), c=torch.ones(2), batch_size=[2])
>>> sorted(td.add(other, default="intersection").keys())
['b']
>>> sorted(td.add(other, default=torch.zeros(())).keys())
['a', 'b', 'c']

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