jax.numpy.cumulative_sum

jax.numpy.cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False)

沿数组轴的累积和。

JAX 实现的 numpy.cumulative_sum().

参数:

• x (ArrayLike) – N 维数组

• axis (int | None | None) – 要累积的整数轴。如果 x 是一个维度，则此参数是可选的。

• dtype (DTypeLike | None | None) – 输出的可选 dtype。

• include_initial (bool) – 如果为 True，则在累积和中包含初始值。默认为 False。

返回值:

包含累积值的数组。

返回类型:

数组

另请参见

• jax.numpy.cumsum(): 累积和的替代 API。

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

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

示例

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.cumulative_sum(x, axis=1)
Array([[ 1,  3,  6],
       [ 4,  9, 15]], dtype=int32)
>>> jnp.cumulative_sum(x, axis=1, include_initial=True)
Array([[ 0,  1,  3,  6],
       [ 0,  4,  9, 15]], dtype=int32)