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 (Optional[Callable[[Any], bool]]) –

Return type:

List[Any]