jax.tree_util.tree_leaves

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

Gets the leaves of a pytree.

Parameters:

• tree (Any) –

• is_leaf (Callable[[Any], bool] | None) –

Return type:

list[Leaf]