jax.tree_util.tree_reduce_associative#

jax.tree_util.tree_reduce_associative(operation, tree, *, identity=<jax._src.tree_util.Unspecified object>, is_leaf=None)[source]#

Alias of jax.tree.reduce_associative().

Parameters:

  • operation (Callable[[T, T], T])

  • tree (Any)

  • identity (T | Unspecified)

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

Return type:

T