TensorDictBase.flatten_keys#

TensorDictBase.flatten_keys(separator: str = '.', inplace: bool = False, is_leaf: Callable[[Type], bool] | None = None) → Self#

Returns a tensordict with the nested entries moved to the top level, under joined keys.

The leaf at ("a", "b", "c") gets the key "a.b.c" with the default separator. The leaves are not copied. The result has the batch size, device and dimension names of the tensordict, and the type of empty(): a TensorDict for most types, a lazy stack for a lazy stack. Two keys that become the same once joined raise a KeyError. unflatten_keys() does the reverse.

It can be used as a context manager: the changes made to the flat tensordict in the with block are written back to the tensordict when the block exits.

Parameters:

separator (str, optional) – the string that joins the parts of a nested key. Defaults to ".".
inplace (bool, optional) – if True, the entries are moved in the tensordict itself, which is returned. Defaults to False.
is_leaf (Callable, optional) – a function that takes a class and returns True if values of that class count as leaves. Here it stops the flattening: a nested tensordict whose class is a leaf is kept as one value under its key. Defaults to None, which uses is_leaf_nontensor(): tensors and non-tensor data are leaves, tensordicts and tensorclasses are flattened.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 1, ("b", "c"): 2, ("e", "f", "g"): 3})
>>> list(td.flatten_keys().keys())
['a', 'b.c', 'e.f.g']
>>> list(td.flatten_keys(separator="/").keys())
['a', 'b/c', 'e/f/g']
>>> with td.flatten_keys() as flat:
...     flat["b.c"] += 1
>>> td["b", "c"]
tensor(3)

With unflatten_keys(), it converts a state dict to a nested tensordict and back:

>>> model = torch.nn.Sequential(torch.nn.Linear(3, 4))
>>> params = TensorDict(model.state_dict()).unflatten_keys(".")
>>> sorted(params.keys(include_nested=True, leaves_only=True))
[('0', 'bias'), ('0', 'weight')]
>>> _ = model.load_state_dict(params.flatten_keys(".").to_dict())