jax.tree_util.tree_reduce

jax.tree_util.tree_reduce(function, tree, initializer=<jax._src.tree_util.Unspecified object>, is_leaf=None)

Alias of jax.tree.reduce().

Parameters:

• function (Callable[[T, Any], T])

• tree (Any)

• initializer (T | Unspecified)

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

Return type:

T