jax.lax.reduce_prod

jax.lax.reduce_prod(operand, axes)

计算一个或多个数组轴上的元素乘积。

参数:

• operand (ArrayLike) – 要在其上求和的数组。必须具有数值 dtype。

• axes (Sequence[int]) – 指定要在其上求和的轴的零个或多个唯一整数的序列。每个条目必须满足 0 <= axis < operand.ndim。

返回:

与 operand 具有相同 dtype 的数组，其形状对应于删除了 axes 的 operand.shape 的维度。

返回类型:

Array

笔记

与 jax.numpy.prod() 不同，jax.lax.reduce_prod() 不会为了累积而向上转换窄位宽类型，因此 8 位或 16 位类型的乘积可能会受到舍入误差的影响。

另请参阅

• jax.numpy.prod()：更灵活的 NumPy 风格的乘积 API，围绕 jax.lax.reduce_prod() 构建。

• 其他低级 jax.lax 归约运算符：jax.lax.reduce_sum(), jax.lax.reduce_max(), jax.lax.reduce_min(), jax.lax.reduce_and(), jax.lax.reduce_or(), jax.lax.reduce_xor()。