Skip to main content
Ctrl+K

jax.tree_util.tree_leaves#

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

Gets the leaves of a pytree.

Parameters:
  • tree (Any) –

  • is_leaf (Optional[Callable[[Any], bool]]) –

Return type:

list[Any]

previous

jax.tree_util.tree_flatten_with_path

next

jax.tree_util.tree_leaves_with_path