jax.tree_util.tree_reduce#

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

jax.tree_util.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

jax.tree_util.tree_reduce