reference/generated/tensordict.TensorDictBase.update

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.update

TensorDictBase.update(input_dict_or_td: T | Mapping[Any, Any] | None = None, clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None, is_leaf: Callable[[Type], bool] | None = None, update_batch_size: bool = False, ignore_lock: bool = False, **kwargs) → Self

Writes the entries of a dictionary or another tensordict into the tensordict, and returns self.

Each entry is written as with set(): by reference, unless clone=True. Nested tensordicts are merged, not replaced: the entries of the tensordict that the input does not have are kept. Keyword arguments are top-level entries: td.update(a=x) writes x at "a", and wins over an entry "a" of the input. A key named like an argument of update (clone, inplace…) must go in the input dictionary.

Warning

update is not atomic. If one entry fails (for example, its shape does not match the batch size), the entries before it are already written. Do not catch the exception and continue with the same tensordict. Update a copy, or make sure that the input is valid first.

Parameters:

• input_dict_or_td (TensorDictBase or dict, optional) – the entries to write. A dictionary can be nested and have tuple keys. If None, only the keyword arguments are written. Defaults to None.

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

• inplace (bool, optional) – if True, a value whose key exists is copied into the existing tensor, which keeps its storage. A new key is added. Defaults to False. See In-place operations.

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 written, without the cost of td.update(other.select(*keys_to_update)). Defaults to None, which writes every entry.

• is_leaf (Callable[[type], bool], optional) – tells which types of values are written whole. The other values are merged key by key. Defaults to None, which is is_leaf_nontensor(). See also default_is_leaf().

• update_batch_size (bool, optional) – if True, sets the batch size of the tensordict to that of the input when they differ. A lazy stack is then refilled with copies of the tensordicts of the input. It needs keys_to_update=None and inplace=False, and the keys of the tensordict must all be in the input. A batch size that is a prefix of the input’s ([] for an input of batch size [2]) is kept, with a FutureWarning: starting with TensorDict 0.17, the default will change to replacing it as well. A batch size whose leading dims are those of the input ([2, 3] for [2]) is kept. Defaults to False.

• ignore_lock (bool, optional) – deprecated: True warns, and the argument will be removed in TensorDict 0.17. On a locked TensorDict, it does not lift the lock: a new or replaced entry still raises RuntimeError (#2260). Use inplace=True, or unlock the tensordict first, as in with td.unlock_(): td.update(other). Defaults to False.

• **kwargs – entries to write at the top level, by key.

Returns:

self.

Raises:

RuntimeError – if the tensordict is locked (see Locking and temporary changes) and the call adds or replaces an entry, or if the leading dimensions of a value do not match the batch size.

Note

With update_batch_size=True, a lazy stack of N tensordicts updated with a lazy stack of M > N tensordicts along the stack dimension gets copies of the M - N extra tensordicts appended, so update can grow a lazy stack. Without it, the extra tensordicts are left out.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.zeros(3), batch_size=[3])
>>> other = TensorDict(a=torch.ones(3), b=torch.ones(3, 4), batch_size=[3])
>>> td = td.update(other)
>>> td["a"] is other["a"]  # written by reference unless clone=True
True
>>> td = td.update(c=torch.zeros(3))  # keyword form
>>> sorted(td.keys())
['a', 'b', 'c']

Nested tensordicts are merged:

>>> td = td.update({"nested": {"x": torch.zeros(3)}})
>>> td = td.update({"nested": {"y": torch.ones(3)}})
>>> sorted(td["nested"].keys())
['x', 'y']

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

Tutorials that use update

Getting, setting, and nesting keys

Getting, setting, and nesting keys