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.

Parameters:

  • operands (Any)

  • init_values (Any)

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

  • dimensions (Sequence[int])

Return type:

Any