jax.tree.flatten#

jax.tree.flatten(tree, is_leaf=None)[source]#

Flattens a pytree.

The flattening order (i.e. the order of elements in the output list) is deterministic, corresponding to a left-to-right depth-first tree traversal.

Parameters:

  • tree (Any) – a pytree to flatten.

  • is_leaf (Callable[[Any], bool] | None | None) – an optionally specified function that will be called at each flattening step. It should return a boolean, with true stopping the traversal and the whole subtree being treated as a leaf, and false indicating the flattening should traverse the current object.

Returns:

A pair where the first element is a list of leaf values and the second element is a treedef representing the structure of the flattened tree.

Return type:

tuple[list[tree_util.Leaf], tree_util.PyTreeDef]

Examples

>>> import jax
>>> vals, treedef = jax.tree.flatten([1, (2, 3), [4, 5]])
>>> vals
[1, 2, 3, 4, 5]
>>> treedef
PyTreeDef([*, (*, *), [*, *]])

See also