Add cumulative_sum wrapper for the numpy-likes

Author: asmeurer

array_api_compat/common/_aliases.py

@@ -264,6 +264,36 @@ def var(
 ) -> ndarray:
     return xp.var(x, axis=axis, ddof=correction, keepdims=keepdims, **kwargs)
 
+# cumulative_sum is renamed from cumsum, and adds the include_initial keyword
+# argument
+
+def cumulative_sum(
+    x: ndarray,
+    /,
+    xp,
+    *,
+    axis: Optional[int] = None,
+    dtype: Optional[Dtype] = None,
+    include_initial: bool = False,
+    **kwargs
+) -> ndarray:
+    # TODO: The standard is not clear about what should happen when x.ndim == 0.
+    if axis is None:
+        if x.ndim > 1:
+            raise ValueError("axis must be specified in cumulative_sum for more than one dimension")
+        axis = 0
+
+    res = xp.cumsum(x, axis=axis, dtype=dtype, **kwargs)
+
+    # np.cumsum does not support include_initial
+    if include_initial:
+        initial_shape = list(x.shape)
+        initial_shape[axis] = 1
+        res = xp.concatenate(
+            [xp.zeros_like(res, shape=initial_shape), 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.
@@ -502,6 +532,7 @@ def unstack(x: ndarray, /, xp, *, axis: int = 0) -> Tuple[ndarray, ...]:
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
-           'astype', 'std', 'var', 'clip', 'permute_dims', 'reshape',
-           'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc', 'matmul',
-           'matrix_transpose', 'tensordot', 'vecdot', 'isdtype', 'unstack']
+           'astype', 'std', 'var', 'cumulative_sum', 'clip', 'permute_dims',
+           'reshape', 'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc',
+           'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype',
+           'unstack']

array_api_compat/cupy/_aliases.py

@@ -49,6 +49,7 @@
 astype = _aliases.astype
 std = get_xp(cp)(_aliases.std)
 var = get_xp(cp)(_aliases.var)
+cumulative_sum = get_xp(cp)(_aliases.cumulative_sum)
 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

@@ -91,6 +91,7 @@ def _dask_arange(
 permute_dims = get_xp(da)(_aliases.permute_dims)
 std = get_xp(da)(_aliases.std)
 var = get_xp(da)(_aliases.var)
+cumulative_sum = get_xp(da)(_aliases.cumulative_sum)
 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

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