jax.tree_util.tree_flatten

jax.tree_util.tree_flatten(tree, is_leaf=None)

Flattens a pytree.

Parameters

• tree – a pytree to flatten.

• is_leaf (Optional[Callable[[Any], bool]]) – an optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf.

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.