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。

返回类型:

Array

Notes

与 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()。