TensorDictBase.set#

TensorDictBase.set(key: NestedKey, item: Tensor | TensorCollection, inplace: bool = False, *, non_blocking: bool = False, **kwargs: Any) → Self#

Sets an entry, and returns the tensordict.

td.set(key, value) is the same as td[key] = value. By default, the value replaces the existing entry (no data is copied). A tuple key sets a nested entry, and creates the missing sub-tensordicts. The leading dimensions of the value must match batch_size. Arrays and numbers are converted to tensors. Strings and other objects are stored as NonTensorData.

A bool, int or float, or a NumPy scalar of one of these kinds, written over an existing NonTensorData entry is stored as a 0-d tensor, and the write emits a FutureWarning. Starting with TensorDict 0.17, the default will change to storing it as NonTensorData, like the entry it replaces.

Parameters:

key (NestedKey) – the key to set: a string, or a tuple of strings for a nested entry. See Keys, nested keys, and locking.
item (Tensor, TensorDictBase or any object) – the value to store.
inplace (bool, optional) – if True and the key exists, the value is copied into the existing tensor, which keeps its storage. If the key does not exist, the entry is added. For a write that raises on a missing key, use set_(). 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.
**kwargs – ignored, with a FutureWarning. Starting with TensorDict 0.17, the default will change to raising a TypeError for them, as set_() does.

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 the value do not match the batch size.
ValueError – if inplace=True and the value does not fit the shape of the existing entry. The error is also a RuntimeError, the type that update_() raises.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({}, batch_size=[3, 4])
>>> td = td.set("x", torch.randn(3, 4))
>>> y = torch.randn(3, 4, 5)
>>> td = td.set("y", y, inplace=True)  # works, even if 'y' is not present yet
>>> td = td.set("y", torch.zeros_like(y), inplace=True)
>>> assert (y==0).all()  # y values are overwritten
>>> td = td.set(("nested", "z"), torch.zeros(3, 4, 2))  # creates td["nested"]
>>> td["nested"].batch_size
torch.Size([3, 4])
>>> td.set("y", torch.ones(5), inplace=True)  # the shape does not match
Traceback (most recent call last):
...
RuntimeError: batch dimension mismatch, got self.batch_size=torch.Size([3, 4]) and value.shape=torch.Size([5]) ...