implement dpnp.cumulative_sum and dpnp.cumulative_prod (#2171)

* implement dpnp.cumulative_sum and dpnp.cumulative_prod

* add requires numpy-2.0.0

* address comments

* update test_sycl_queue.py to include cumulative_sum and cumulative_prod only if NumPy > = 2.1.0

Author: vtavana

dpnp/dpnp_iface_mathematical.py

@@ -94,6 +94,8 @@
     "cross",
     "cumprod",
     "cumsum",
+    "cumulative_prod",
+    "cumulative_sum",
     "diff",
     "divide",
     "ediff1d",
@@ -1062,6 +1064,8 @@ def cumprod(a, axis=None, dtype=None, out=None):
 
     See Also
     --------
+    :obj:`dpnp.cumulative_prod` : Array API compatible alternative for
+                                :obj:`dpnp.cumprod`.
     :obj:`dpnp.prod` : Product array elements.
 
     Examples
@@ -1143,6 +1147,8 @@ def cumsum(a, axis=None, dtype=None, out=None):
 
     See Also
     --------
+    :obj:`dpnp.cumulative_sum` : Array API compatible alternative for
+                                :obj:`dpnp.cumsum`.
     :obj:`dpnp.sum` : Sum array elements.
     :obj:`dpnp.trapezoid` : Integration of array values using composite
                             trapezoidal rule.
@@ -1157,8 +1163,8 @@ def cumsum(a, axis=None, dtype=None, out=None):
            [4, 5, 6]])
     >>> np.cumsum(a)
     array([ 1,  3,  6, 10, 15, 21])
-    >>> np.cumsum(a, dtype=float)     # specifies type of output value(s)
-    array([ 1.,  3.,  6., 10., 15., 21.])
+    >>> np.cumsum(a, dtype=np.float32)     # specifies type of output value(s)
+    array([ 1.,  3.,  6., 10., 15., 21.], dtype=np.float32)
 
     >>> np.cumsum(a, axis=0)     # sum over rows for each of the 3 columns
     array([[1, 2, 3],
@@ -1169,8 +1175,8 @@ def cumsum(a, axis=None, dtype=None, out=None):
 
     ``cumsum(b)[-1]`` may not be equal to ``sum(b)``
 
-    >>> b = np.array([1, 2e-9, 3e-9] * 10000)
-    >>> b.cumsum().dtype == b.sum().dtype == np.float64
+    >>> b = np.array([1, 2e-7, 3e-7] * 100000, dtype=np.float32)
+    >>> b.cumsum().dtype == b.sum().dtype == np.float32
     True
     >>> b.cumsum()[-1] == b.sum()
     array(False)
@@ -1194,6 +1200,184 @@ def cumsum(a, axis=None, dtype=None, out=None):
     )
 
 
+def cumulative_prod(
+    x, /, *, axis=None, dtype=None, out=None, include_initial=False
+):
+    """
+    Return the cumulative product of elements along a given axis.
+
+    This function is an Array API compatible alternative to :obj:`dpnp.cumprod`.
+
+    For full documentation refer to :obj:`numpy.cumulative_prod`.
+
+    Parameters
+    ----------
+    x : {dpnp.ndarray, usm_ndarray}
+        Input array.
+    axis : {None, int}, optional
+        Axis along which the cumulative product is computed. The default value
+        is only allowed for one-dimensional arrays. For arrays with more than
+        one dimension `axis` is required.
+        Default: ``None``.
+    dtype : {None, dtype}, optional
+        Type of the returned array and of the accumulator in which the elements
+        are summed. If `dtype` is not specified, it defaults to the dtype of
+        `x`, unless `x` has an integer dtype with a precision less than that of
+        the default platform integer. In that case, the default platform
+        integer is used.
+        Default: ``None``.
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
+        Alternative output array in which to place the result. It must have the
+        same shape and buffer length as the expected output but the type will
+        be cast if necessary.
+        Default: ``None``.
+    include_initial : bool, optional
+        Boolean indicating whether to include the initial value (ones) as
+        the first value in the output. With ``include_initial=True``
+        the shape of the output is different than the shape of the input.
+        Default: ``False``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        A new array holding the result is returned unless `out` is specified,
+        in which case a reference to `out` is returned. The
+        result has the same shape as `x` if ``include_initial=False``.
+
+    See Also
+    --------
+    :obj:`dpnp.prod` : Product array elements.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.array([1, 2, 3])
+    >>> np.cumulative_prod(a) # intermediate results 1, 1*2
+    ...                       # total product 1*2*3 = 6
+    array([1, 2, 6])
+    >>> a = np.array([1, 2, 3, 4, 5, 6])
+    >>> np.cumulative_prod(a, dtype=np.float32) # specify type of output
+    array([  1.,   2.,   6.,  24., 120., 720.], dtype=float32)
+
+    The cumulative product for each column (i.e., over the rows) of `b`:
+
+    >>> b = np.array([[1, 2, 3], [4, 5, 6]])
+    >>> np.cumulative_prod(b, axis=0)
+    array([[ 1,  2,  3],
+           [ 4, 10, 18]])
+
+    The cumulative product for each row (i.e. over the columns) of `b`:
+
+    >>> np.cumulative_prod(b, axis=1)
+    array([[  1,   2,   6],
+           [  4,  20, 120]])
+
+    """
+
+    return dpnp_wrap_reduction_call(
+        x,
+        out,
+        dpt.cumulative_prod,
+        _get_reduction_res_dt,
+        dpnp.get_usm_ndarray(x),
+        axis=axis,
+        dtype=dtype,
+        include_initial=include_initial,
+    )
+
+
+def cumulative_sum(
+    x, /, *, axis=None, dtype=None, out=None, include_initial=False
+):
+    """
+    Return the cumulative sum of the elements along a given axis.
+
+    This function is an Array API compatible alternative to :obj:`dpnp.cumsum`.
+
+    For full documentation refer to :obj:`numpy.cumulative_sum`.
+
+    Parameters
+    ----------
+    x : {dpnp.ndarray, usm_ndarray}
+        Input array.
+    axis : {None, int}, optional
+        Axis along which the cumulative sum is computed. The default value
+        is only allowed for one-dimensional arrays. For arrays with more than
+        one dimension `axis` is required.
+        Default: ``None``.
+    dtype : {None, dtype}, optional
+        Type of the returned array and of the accumulator in which the elements
+        are summed. If `dtype` is not specified, it defaults to the dtype of
+        `x`, unless `x` has an integer dtype with a precision less than that of
+        the default platform integer. In that case, the default platform
+        integer is used.
+        Default: ``None``.
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
+        Alternative output array in which to place the result. It must have the
+        same shape and buffer length as the expected output but the type will
+        be cast if necessary.
+        Default: ``None``.
+    include_initial : bool, optional
+        Boolean indicating whether to include the initial value (ones) as
+        the first value in the output. With ``include_initial=True``
+        the shape of the output is different than the shape of the input.
+        Default: ``False``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        A new array holding the result is returned unless `out` is specified,
+        in which case a reference to `out` is returned. The
+        result has the same shape as `x` if ``include_initial=False``.
+
+    See Also
+    --------
+    :obj:`dpnp.sum` : Sum array elements.
+    :obj:`dpnp.trapezoid` : Integration of array values using composite
+                            trapezoidal rule.
+    :obj:`dpnp.diff` : Calculate the n-th discrete difference along given axis.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.array([1, 2, 3, 4, 5, 6])
+    >>> a
+    array([1, 2, 3, 4, 5, 6])
+    >>> np.cumulative_sum(a)
+    array([ 1,  3,  6, 10, 15, 21])
+    >>> np.cumulative_sum(a, dtype=np.float32)   # specifies output dtype
+    array([ 1.,  3.,  6., 10., 15., 21.], dtype=np.float32)
+
+    >>> b = np.array([[1, 2, 3], [4, 5, 6]])
+    >>> np.cumulative_sum(b, axis=0)  # sum over rows for each of the 3 columns
+    array([[1, 2, 3],
+           [5, 7, 9]])
+    >>> np.cumulative_sum(b, axis=1)  # sum over columns for each of the 2 rows
+    array([[ 1,  3,  6],
+           [ 4,  9, 15]])
+
+    ``cumulative_sum(c)[-1]`` may not be equal to ``sum(c)``
+
+    >>> c = np.array([1, 2e-7, 3e-7] * 100000, dtype=np.float32)
+    >>> np.cumulative_sum(c).dtype == c.sum().dtype == np.float32
+    True
+    >>> np.cumulative_sum(c)[-1] == c.sum()
+    array(False)
+
+    """
+
+    return dpnp_wrap_reduction_call(
+        x,
+        out,
+        dpt.cumulative_sum,
+        _get_reduction_res_dt,
+        dpnp.get_usm_ndarray(x),
+        axis=axis,
+        dtype=dtype,
+        include_initial=include_initial,
+    )
+
+
 def diff(a, n=1, axis=-1, prepend=None, append=None):
     """
     Calculate the n-th discrete difference along the given axis.

tests/test_mathematical.py

@@ -424,6 +424,39 @@ def test_out_dtype(self, arr_dt, out_dt, dtype):
         assert_array_equal(expected, result)
         assert result is iout
 
+    @testing.with_requires("numpy>=2.1.0")
+    def test_include_initial(self):
+        a = numpy.arange(8).reshape(2, 2, 2)
+        ia = dpnp.array(a)
+
+        expected = numpy.cumulative_prod(a, axis=1, include_initial=True)
+        result = dpnp.cumulative_prod(ia, axis=1, include_initial=True)
+        assert_array_equal(result, expected)
+
+        expected = numpy.cumulative_prod(a, axis=0, include_initial=True)
+        result = dpnp.cumulative_prod(ia, axis=0, include_initial=True)
+        assert_array_equal(result, expected)
+
+        a = numpy.arange(1, 5).reshape(2, 2)
+        ia = dpnp.array(a)
+        out = numpy.zeros((3, 2), dtype=numpy.float32)
+        out_dp = dpnp.array(out)
+
+        expected = numpy.cumulative_prod(
+            a, axis=0, out=out, include_initial=True
+        )
+        result = dpnp.cumulative_prod(
+            ia, axis=0, out=out_dp, include_initial=True
+        )
+        assert result is out_dp
+        assert_array_equal(result, expected)
+
+        a = numpy.array([2, 2])
+        ia = dpnp.array(a)
+        expected = numpy.cumulative_prod(a, include_initial=True)
+        result = dpnp.cumulative_prod(ia, include_initial=True)
+        assert_array_equal(result, expected)
+
 
 class TestCumSum:
     @pytest.mark.parametrize(
@@ -495,6 +528,39 @@ def test_out_dtype(self, arr_dt, out_dt, dtype):
         assert_array_equal(expected, result)
         assert result is iout
 
+    @testing.with_requires("numpy>=2.1.0")
+    def test_include_initial(self):
+        a = numpy.arange(8).reshape(2, 2, 2)
+        ia = dpnp.array(a)
+
+        expected = numpy.cumulative_sum(a, axis=1, include_initial=True)
+        result = dpnp.cumulative_sum(ia, axis=1, include_initial=True)
+        assert_array_equal(result, expected)
+
+        expected = numpy.cumulative_sum(a, axis=0, include_initial=True)
+        result = dpnp.cumulative_sum(ia, axis=0, include_initial=True)
+        assert_array_equal(result, expected)
+
+        a = numpy.arange(1, 5).reshape(2, 2)
+        ia = dpnp.array(a)
+        out = numpy.zeros((3, 2), dtype=numpy.float32)
+        out_dp = dpnp.array(out)
+
+        expected = numpy.cumulative_sum(
+            a, axis=0, out=out, include_initial=True
+        )
+        result = dpnp.cumulative_sum(
+            ia, axis=0, out=out_dp, include_initial=True
+        )
+        assert result is out_dp
+        assert_array_equal(result, expected)
+
+        a = numpy.array([2, 2])
+        ia = dpnp.array(a)
+        expected = numpy.cumulative_sum(a, include_initial=True)
+        result = dpnp.cumulative_sum(ia, include_initial=True)
+        assert_array_equal(result, expected)
+
 
 class TestDiff:
     @pytest.mark.parametrize("n", list(range(0, 3)))
@@ -1927,7 +1993,7 @@ def test_zero(self, dt):
         expected = numpy.sinc(a)
         assert_dtype_allclose(result, expected)
 
-    # TODO: add a proper NumPY version once resolved
+    # TODO: add a proper NumPy version once resolved
     @testing.with_requires("numpy>=2.0.0")
     def test_zero_fp16(self):
         a = numpy.array([0.0], dtype=numpy.float16)

tests/test_sycl_queue.py

@@ -464,6 +464,8 @@ def test_meshgrid(device):
         pytest.param("count_nonzero", [3, 0, 2, -1.2]),
         pytest.param("cumprod", [[1, 2, 3], [4, 5, 6]]),
         pytest.param("cumsum", [[1, 2, 3], [4, 5, 6]]),
+        pytest.param("cumulative_prod", [1, 2, 3, 4, 5, 6]),
+        pytest.param("cumulative_sum", [1, 2, 3, 4, 5, 6]),
         pytest.param("degrees", [numpy.pi, numpy.pi / 2, 0]),
         pytest.param("diagonal", [[[1, 2], [3, 4]]]),
         pytest.param("diff", [1.0, 2.0, 4.0, 7.0, 0.0]),
@@ -556,6 +558,14 @@ def test_1in_1out(func, data, device):
         # `trapezoid` is available from NumPy 2.0
         func = "trapz"
 
+    if (
+        func in ["cumulative_prod", "cumulative_sum"]
+        and numpy.lib.NumpyVersion(numpy.__version__) < "2.1.0"
+    ):
+        pytest.skip(
+            "cumulative_prod and cumulative_sum are available from NumPy 2.1"
+        )
+
     x_orig = dpnp.asnumpy(x)
     expected = getattr(numpy, func)(x_orig)
     assert_dtype_allclose(result, expected)

tests/test_usm_type.py

@@ -585,6 +585,8 @@ def test_norm(usm_type, ord, axis):
         pytest.param("cumlogsumexp", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("cumprod", [[1, 2, 3], [4, 5, 6]]),
         pytest.param("cumsum", [[1, 2, 3], [4, 5, 6]]),
+        pytest.param("cumulative_prod", [1, 2, 3, 4, 5, 6]),
+        pytest.param("cumulative_sum", [1, 2, 3, 4, 5, 6]),
         pytest.param("degrees", [numpy.pi, numpy.pi / 2, 0]),
         pytest.param("diagonal", [[[1, 2], [3, 4]]]),
         pytest.param("diff", [1.0, 2.0, 4.0, 7.0, 0.0]),