ENH: add cumulative_prod

Author: ev-br

array_api_compat/common/_aliases.py

@@ -297,6 +297,36 @@ def cumulative_sum(
         )
     return res
 
+
+def cumulative_prod(
+    x: ndarray,
+    /,
+    xp,
+    *,
+    axis: Optional[int] = None,
+    dtype: Optional[Dtype] = None,
+    include_initial: bool = False,
+    **kwargs
+) -> ndarray:
+    wrapped_xp = array_namespace(x)
+
+    if axis is None:
+        if x.ndim > 1:
+            raise ValueError("axis must be specified in cumulative_prod for more than one dimension")
+        axis = 0
+
+    res = xp.cumprod(x, axis=axis, dtype=dtype, **kwargs)
+
+    # np.cumprod does not support include_initial
+    if include_initial:
+        initial_shape = list(x.shape)
+        initial_shape[axis] = 1
+        res = xp.concatenate(
+            [wrapped_xp.ones(shape=initial_shape, dtype=res.dtype, device=device(res)), res],
+            axis=axis,
+        )
+    return res
+
 # The min and max argument names in clip are different and not optional in numpy, and type
 # promotion behavior is different.
 def clip(
@@ -549,7 +579,7 @@ def sign(x: ndarray, /, xp, **kwargs) -> ndarray:
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
-           'astype', 'std', 'var', 'cumulative_sum', 'clip', 'permute_dims',
+           'astype', 'std', 'var', 'cumulative_sum', 'cumulative_prod', 'clip', 'permute_dims',
            'reshape', 'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc',
            'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype',
            'unstack', 'sign']

array_api_compat/cupy/_aliases.py

@@ -50,6 +50,7 @@
 std = get_xp(cp)(_aliases.std)
 var = get_xp(cp)(_aliases.var)
 cumulative_sum = get_xp(cp)(_aliases.cumulative_sum)
+cumulative_prod = get_xp(cp)(_aliases.cumulative_prod)
 clip = get_xp(cp)(_aliases.clip)
 permute_dims = get_xp(cp)(_aliases.permute_dims)
 reshape = get_xp(cp)(_aliases.reshape)

array_api_compat/dask/array/_aliases.py

@@ -86,6 +86,7 @@ def _dask_arange(
 std = get_xp(da)(_aliases.std)
 var = get_xp(da)(_aliases.var)
 cumulative_sum = get_xp(da)(_aliases.cumulative_sum)
+cumulative_prod = get_xp(da)(_aliases.cumulative_prod)
 empty = get_xp(da)(_aliases.empty)
 empty_like = get_xp(da)(_aliases.empty_like)
 full = get_xp(da)(_aliases.full)

array_api_compat/numpy/_aliases.py

@@ -50,6 +50,7 @@
 std = get_xp(np)(_aliases.std)
 var = get_xp(np)(_aliases.var)
 cumulative_sum = get_xp(np)(_aliases.cumulative_sum)
+cumulative_prod = get_xp(np)(_aliases.cumulative_prod)
 clip = get_xp(np)(_aliases.clip)
 permute_dims = get_xp(np)(_aliases.permute_dims)
 reshape = get_xp(np)(_aliases.reshape)

array_api_compat/torch/_aliases.py

@@ -204,6 +204,7 @@ def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keep
 clip = get_xp(torch)(_aliases_clip)
 unstack = get_xp(torch)(_aliases_unstack)
 cumulative_sum = get_xp(torch)(_aliases_cumulative_sum)
+cumulative_prod = get_xp(torch)(_aliases_cumulative_prod)
 
 # torch.sort also returns a tuple
 # https://github.com/pytorch/pytorch/issues/70921