flat_mapping

Contents

flat_mapping#

class brainstate.util.flat_mapping(xs, /, *, keep_empty_nodes=False, is_leaf=<function <lambda>>, sep=None)[source]#

Flatten a nested mapping into a flat mapping with tuple or string keys.

The nested keys are flattened to a tuple path. For example, {'a': {'b': 1}} becomes {('a', 'b'): 1}. See nest_mapping() on how to restore the nested structure.

Parameters:

  • xs (Mapping[Any, Any]) – A nested mapping to flatten.

  • keep_empty_nodes (bool) – If True, replaces empty mappings with empty_node sentinel. Otherwise, empty mappings are omitted from the result.

  • is_leaf (Callable[[Tuple[Any, ...], Mapping[Any, Any]], bool] | None) – Optional function that takes (path, mapping) and returns True if the mapping should be treated as a leaf (i.e., not flattened further). Defaults to treating all mappings as non-leaves.

  • sep (str | None) – If specified, keys in the returned mapping will be sep-joined strings instead of tuples. For example, with sep='/', ('a', 'b') becomes 'a/b'.

Returns:

A flattened mapping where nested keys are converted to tuples or strings.

Return type:

FlattedDict

Example

>>> xs = {'foo': 1, 'bar': {'a': 2, 'b': {}}}
>>> flat_xs = flat_mapping(xs)
>>> flat_xs
FlattedDict({('foo',): 1, ('bar', 'a'): 2})

>>> # With separator
>>> flat_mapping(xs, sep='/')
FlattedDict({'foo': 1, 'bar/a': 2})

>>> # Keep empty nodes
>>> flat_mapping(xs, keep_empty_nodes=True)
FlattedDict({('foo',): 1, ('bar', 'a'): 2, ('bar', 'b'): _EmptyNode()})

Note

Empty mappings are ignored by default and will not be restored by nest_mapping() unless keep_empty_nodes=True.

Contents