jax.tree_util.tree_leaves#

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

Gets the leaves of a pytree.

Parameters
Return type

List[Any]