TensorDictBase.select#

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

Selects the keys of the tensordict and returns a new tensordict with only the selected keys.

The values are not copied. The result shares its tensors with the original, so an in-place change to a tensor in one tensordict shows in the other. A nested key selects one entry of a sub-tensordict, and a key of a sub-tensordict selects the whole sub-tensordict. See Keys, nested keys, and locking.

Parameters:

*keys (NestedKey) – keys to select.
inplace (bool) – if True, the other keys are removed from this tensordict, and this tensordict is returned. Default is False.
strict (bool, optional) – if True, a missing key raises KeyError. If False, missing keys are ignored. Default: True.
as_tensordict (bool, optional) – if True, the result will be a plain TensorDict even when called on a TensorClass instance. This avoids the TensorClass wrapper that fills unselected fields with None. For TensorDictBase subclasses, this is a no-op. Default: False.

Returns:

A new tensordict (or the same if inplace=True) with the selected keys only.

Note

To select keys in a tensordict and return a version of this tensordict deprived of these keys, see the split_keys() method.

Examples

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, [])
>>> td.select("a", ("b", "c"))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("a", "b")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("this key does not exist", strict=False)
TensorDict(
    fields={
    },
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)