jax.tree_util.tree_reduce_associative

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

jax.tree.reduce_associative() 的别名。

参数:

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

• tree (Any)

• identity (T | Unspecified)

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

返回类型:

T