mlx.utils.tree_flatten

Contents

mlx.utils.tree_flatten#

tree_flatten(tree: Any, prefix: str = '', is_leaf: Callable | None = None, destination: List[Tuple[str, Any]] | Dict[str, Any] | None = None) → List[Tuple[str, Any]] | Dict[str, Any]#

Flattens a Python tree to a list of key, value tuples.

The keys are using the dot notation to define trees of arbitrary depth and complexity.

from mlx.utils import tree_flatten

print(tree_flatten([[[0]]]))
# [("0.0.0", 0)]

print(tree_flatten([[[0]]], prefix=".hello"))
# [("hello.0.0.0", 0)]

tree_flatten({"a": {"b": 1}}, destination={})
{"a.b": 1}

Note

Dictionaries should have keys that are valid Python identifiers.

Parameters:

  • tree (Any) – The Python tree to be flattened.

  • prefix (str) – A prefix to use for the keys. The first character is always discarded.

  • is_leaf (callable) – An optional callable that returns True if the passed object is considered a leaf or False otherwise.

  • destination (list or dict, optional) – A list or dictionary to store the flattened tree. If None an empty list will be used. Default: None.

Returns:

The flat representation of

the Python tree.

Return type:

Union[List[Tuple[str, Any]], Dict[str, Any]]

Contents