jax.numpy.cumprod

目录

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) – 可选地指定输出的数据类型。如果未指定,则输出数据类型将与输入数据类型匹配。

  • out (None) – JAX 未使用

返回:

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

返回类型:

Array

参见

示例

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