jax.numpy.cumprod

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

沿轴的元素的累积乘积。

JAX 对 numpy.cumprod() 的实现。

参数:

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

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

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

• out (None) – JAX 未使用。

返回:

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

返回类型:

数组

另请参阅

• jax.numpy.multiply.accumulate(): 通过 ufunc 方法进行累积乘积。

• jax.numpy.nancumprod(): 忽略 NaN 值的累积乘积。

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

示例

>>> x = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.cumprod(x)  # flattened cumulative product
Array([  1,   2,   6,  24, 120, 720], dtype=int32)
>>> jnp.cumprod(x, axis=1)  # cumulative product along axis 1
Array([[  1,   2,   6],
       [  4,  20, 120]], dtype=int32)