jax.tree.leaves#

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

Gets the leaves of a pytree.

Parameters:

  • tree (Any) – the pytree for which to get the leaves

  • is_leaf (Callable[[Any], bool] | None | None) – an optionally specified function that will be called at each flattening step. It should return a boolean, which indicates whether the flattening should traverse the current object, or if it should be stopped immediately, with the whole subtree being treated as a leaf.

Returns:

a list of tree leaves.

Return type:

leaves

Examples

>>> import jax
>>> jax.tree.leaves([1, (2, 3), [4, 5]])
[1, 2, 3, 4, 5]

See also