jax.lax.reduce_sum

jax.lax.reduce_sum(operand, axes)

计算一个或多个数组轴上的元素的总和。

参数:

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

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

返回:

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

返回类型:

数组

注意事项

与 jax.numpy.sum() 不同，jax.lax.reduce_sum() 不会对窄宽度类型进行向上转型以进行累积，因此 8 位或 16 位类型的总和可能会受到舍入误差的影响。

另请参阅

• jax.numpy.sum(): 更灵活的 NumPy 风格的求和 API，围绕 jax.lax.reduce_sum() 构建。

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