jax.numpy.nancumsum

内容

jax.numpy.nancumsum#

jax.numpy.nancumsum(a, axis=None, dtype=None, out=None)[source]#

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

JAX 实现 numpy.nancumsum().

参数：

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

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

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

  • out (None) – JAX 未使用

返回值：

包含沿给定轴累加的和的数组。

返回类型：

Array

另请参阅

示例

>>> 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)
目录