jax.lax.reduce

jax.lax.reduce(operands, init_values, computation, dimensions)

Wraps XLA’s Reduce operator.

init_values and computation together must form a monoid for correctness. That is init_values must be an identity of computation, and computation must be associative. XLA may exploit both of these properties during code generation; if either is violated the result is undefined.

Parameters:

• operands (Any)

• init_values (Any)

• computation (Callable[[Any, Any], Any])

• dimensions (Sequence[int])

Return type:

Any