jax.tree_util.tree_leaves#

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

Alias of jax.tree.leaves().

Parameters:
  • tree (Any)

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

Return type:

list[Leaf]