jax.lax.reduce¶

jax.lax.reduce(operands, init_values, computation, dimensions)[source]¶

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.