jax.numpy.nancumprod

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

沿轴计算元素的累积乘积，忽略 NaN 值。

JAX 对 numpy.nancumprod() 的实现。

参数:

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

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

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

• out (None) – JAX 未使用

返回:

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

返回类型:

数组

另请参阅

• jax.numpy.cumprod()：不忽略 NaN 值的累积乘积。

• jax.numpy.multiply.accumulate()：通过 ufunc 方法计算累积乘积。

• jax.numpy.prod()：沿轴的乘积

示例

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

标准的累积乘积将传播 NaN 值

>>> jnp.cumprod(x)
Array([ 1.,  2., nan, nan, nan, nan], dtype=float32)

nancumprod() 将忽略 NaN 值，实际上将其替换为 1

>>> jnp.nancumprod(x)
Array([ 1.,  2.,  2.,  8.,  8., 48.], dtype=float32)

沿轴 1 计算累积乘积

>>> jnp.nancumprod(x, axis=1)
Array([[ 1.,  2.,  2.],
       [ 4.,  4., 24.]], dtype=float32)