reference/generated/tensordict.TensorDictBase.keys

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.keys

abstract TensorDictBase.keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Callable[[Type], bool] | None = None, *, sort: bool = False)

Returns a view of the keys of the tensordict.

By default, only the top-level keys are returned. A nested key is a tuple of strings, for example ("nested", "b"). See Keys, nested keys, and locking.

Warning

The view is lazy. If you modify the tensordict after you call keys() and before you iterate over the result, the iteration returns the keys after the change.

Note

On a TensorDict, and on most other types, key in td.keys(include_nested=True, leaves_only=True) checks that key exists, not that it is a leaf: it is True for a sub-tensordict, although iterating the view does not list it. To test that a key is a leaf, call is_leaf(td.entry_class(key)), with default_is_leaf() as the default is_leaf. A nested key is never in a view built with include_nested=False: ("nested", "b") in td.keys() is False. On a LazyStackedTensorDict, a TensorDictStore and a LazyStackedTensorDictStore, it is still True, with a FutureWarning, until TensorDict 0.17.

Parameters:

• include_nested (bool, optional) – if True, the entries of the sub-tensordicts are also returned, with nested keys (tuples of strings). Defaults to False.

• leaves_only (bool, optional) – if True, only leaves are returned, and sub-tensordicts are skipped. Defaults to False.

• is_leaf (callable, optional) – a function that takes a class and returns True if values of that class count as leaves. It only filters the result when leaves_only=True. It does not stop the recursion: with include_nested=True, a sub-tensordict is still searched even if is_leaf returns True for it. By default, tensordicts and tensorclasses are not leaves. See is_leaf_nontensor() and default_is_leaf().

Keyword Arguments:

sort (bool, optional) – if True, the keys are sorted. Nested keys are sorted by their joined name: ("a", "key") sorts as "a.key". Sorting takes extra time on large tensordicts. Defaults to False.

Examples

>>> from tensordict import TensorDict
>>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[])
>>> list(data.keys())
['0', '1']
>>> list(data.keys(leaves_only=True))
['0']
>>> list(data.keys(include_nested=True, leaves_only=True))
['0', ('1', '2')]

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

Tutorials that use keys

Getting, setting, and nesting keys

Getting, setting, and nesting keys