TensorDictBase.select#

TensorDictBase.select(*keys: NestedKey, inplace: bool = False, strict: bool = True, as_tensordict: bool = False) → Self#

Returns a tensordict that holds only the entries of keys.

The values are not copied: the result holds the same tensors as the tensordict, so an in-place change to a tensor shows in both. A nested key such as ("next", "obs") keeps one entry of a nested tensordict, and the key of a nested tensordict keeps all its entries. The result is not locked. See Keys, nested keys, and locking.

Parameters:

*keys (NestedKey) – the keys to keep.

Keyword Arguments:

inplace (bool, optional) – if True, removes the other entries from the tensordict and returns it. Defaults to False.
strict (bool, optional) – if True, a missing key raises a KeyError. If False, missing keys are ignored. Defaults to True.
as_tensordict (bool, optional) – if True, a tensorclass returns a TensorDict instead of an instance of its class, whose other fields are None. Has no effect on other tensordicts. Defaults to False.

Returns:

A tensordict of the same type, or self if inplace=True.

See also

exclude() keeps the other entries, and split_keys() returns both parts.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
...     obs=torch.zeros(3, 4),
...     action=torch.zeros(3, 2),
...     next=TensorDict(obs=torch.zeros(3, 4), reward=torch.zeros(3, 1)),
...     batch_size=[3],
... )
>>> sub = td.select("obs", ("next", "reward"))
>>> list(sub.keys(include_nested=True, leaves_only=True))
['obs', ('next', 'reward')]
>>> td.select("obs")["obs"] is td["obs"]
True
>>> list(td.select("obs", "missing", strict=False).keys())
['obs']