jax.numpy.cumsum

jax.numpy.cumsum(a, axis=None, dtype=None, out=None)

沿轴的元素的累积和。

JAX 对 numpy.cumsum() 的实现。

参数:

• a (ArrayLike) – 要累加的 N 维数组。

• axis (int | None) – 累加的整数轴。如果为 None（默认值），则数组将被展平，并沿展平的轴累加。

• dtype (DTypeLike | None) – 可选地指定输出的数据类型。如果未指定，则输出的数据类型将与输入的数据类型匹配。

• out (None) – JAX 未使用。

返回:

一个包含沿给定轴的累积和的数组。

返回类型:

Array

另请参阅

• jax.numpy.cumulative_sum(): 通过数组 API 标准进行累积求和。

• jax.numpy.add.accumulate(): 通过 ufunc 方法进行累积求和。

• jax.numpy.nancumsum(): 忽略 NaN 值的累积求和。

• jax.numpy.sum(): 沿轴求和

示例

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.cumsum(x)  # flattened cumulative sum
Array([ 1,  3,  6, 10, 15, 21], dtype=int32)
>>> jnp.cumsum(x, axis=1)  # cumulative sum along axis 1
Array([[ 1,  3,  6],
       [ 4,  9, 15]], dtype=int32)