jax.tree_util.tree_flatten#
- jax.tree_util.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 (
Optional[Callable[[Any],bool]]) – 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.
- Return type:
- 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.