jax.tree.reduce#

jax.tree.reduce(function: Callable[[T, Any], T], tree: Any, *, is_leaf: Callable[[Any], bool] | None = None) → T[source]#
jax.tree.reduce(function: Callable[[T, Any], T], tree: Any, initializer: T, is_leaf: Callable[[Any], bool] | None = None) → T

Call reduce() over the leaves of a tree.

Parameters:

  • function – the reduction function

  • tree – the pytree to reduce over

  • initializer – the optional initial value

  • is_leaf – 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:

the reduced value.

Return type:

result

Examples

>>> import jax
>>> import operator
>>> jax.tree.reduce(operator.add, [1, (2, 3), [4, 5, 6]])
21

See also