jax.numpy.nancumsum

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

沿轴的元素累积和，忽略 NaN 值。

JAX 实现的 numpy.nancumsum()。

参数:

• a (类数组) – 要累积的 N 维数组。

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

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

• out (None) – JAX 未使用此参数。

返回值:

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

返回类型:

Array

另请参阅

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

• jax.numpy.cumulative_sum(): 通过数组 API 标准实现的累积和。

• jax.numpy.add.accumulate(): 通过 ufunc 方法实现的累积和。

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

示例

>>> x = jnp.array([[1., 2., jnp.nan],
...                [4., jnp.nan, 6.]])

标准的累积和会传播 NaN 值

>>> jnp.cumsum(x)
Array([ 1.,  3., nan, nan, nan, nan], dtype=float32)

nancumsum() 将忽略 NaN 值，实际上将其替换为零

>>> jnp.nancumsum(x)
Array([ 1.,  3.,  3.,  7.,  7., 13.], dtype=float32)

沿轴 1 的累积和

>>> jnp.nancumsum(x, axis=1)
Array([[ 1.,  3.,  3.],
       [ 4.,  4., 10.]], dtype=float32)