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