jax.lax.reduce

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

包装 XLA 的 Reduce 运算符。

init_values 和 computation 必须共同形成一个 幺半群 才能保证正确性。也就是说，init_values 必须是 computation 的单位元，而 computation 必须是结合的。XLA 可能会在代码生成期间利用这两个特性；如果其中任何一个特性被违反，结果将是未定义的。

参数：

• operands (Any)

• init_values (Any)

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

• dimensions (Sequence[int])

返回值类型：

Any