reference/generated/tensordict.TensorDictBase.update_at_

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.update_at

TensorDictBase.update_at_(input_dict_or_td: dict[str, Tensor | TensorCollection] | T, idx: None | int | slice | str | Tensor | List[Any] | Tuple[Any, ...], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None) → Self

Copies the entries of a dictionary or tensordict into the existing entries at idx.

td.update_at_(other, idx) writes each entry as td.set_at_(key, other[key], idx): the existing tensors keep their storage, so update_at_ works on a locked tensordict. Every key of the input must exist in the tensordict. Tensor and non-tensor leaves are supported. For a faster copy of tensors only, use copy_at_().

Parameters:

• input_dict_or_td (TensorDictBase or dict) – the entries to copy. Their leading dimensions are those of td[idx].

• idx (int, slice, Tensor or tuple) – the index to write at, along the batch dimensions.

• clone (bool, optional) – if True, clones the tensors of the input before copying them. Defaults to False.

Keyword Arguments:

• non_blocking (bool, optional) – if True, copies between devices can run asynchronously with respect to the host. Defaults to False.

• keys_to_update (sequence of NestedKey, optional) – if passed, only these keys are copied. Defaults to None, which copies every entry. If it lists a sub-tensordict key and a key under it, the whole sub-tensordict is copied and the call warns, while select keeps only the deeper keys; starting with TensorDict 0.17, the default will change to copying only the deeper keys, as select does.

Returns:

self.

Raises:

• KeyError – if a key of the input is missing from the tensordict.

• RuntimeError – if a value does not fit the shape of its entry.

Deprecated since version 0.15: The keyword names index (for idx) and input_dict (for input_dict_or_td) warn, and will be removed in TensorDict 0.17.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3, 4, 5), b=torch.zeros(3, 4, 10), batch_size=[3, 4])
>>> other = TensorDict(a=torch.ones(1, 4, 5), b=torch.ones(1, 4, 10), batch_size=[1, 4])
>>> _ = td.update_at_(other, slice(1, 2))
>>> td["a"].sum(dim=(1, 2))
tensor([ 0., 20.,  0.])

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