Add cumulative_sum and cumulative_prod (Array API)

Author: antonwolfy

dpnp/tests/third_party/cupy/math_tests/test_sumprod.py

@@ -1357,3 +1357,130 @@ def test_trapz_1dim_with_x_and_dx(self, xp, dtype):
         a = testing.shaped_arange((5,), xp, dtype)
         x = testing.shaped_arange((5,), xp, dtype)
         return xp.trapezoid(a, x=x, dx=0.1)
+
+
+@testing.with_requires("numpy>=2.1")
+@pytest.mark.parametrize("func", ["cumulative_sum", "cumulative_prod"])
+class TestCumulativeSumProd:
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_1d(self, xp, dtype, func):
+        a = testing.shaped_arange((5,), xp, dtype)
+        return getattr(xp, func)(a)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6, contiguous_check=False)
+    def test_axis(self, xp, dtype, func):
+        a = testing.shaped_arange((3, 4, 5), xp, dtype)
+        return getattr(xp, func)(a, axis=1)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_negative_axis(self, xp, dtype, func):
+        a = testing.shaped_arange((3, 4, 5), xp, dtype)
+        return getattr(xp, func)(a, axis=-1)
+
+    @testing.for_all_dtypes(no_bool=True)
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_dtype(self, xp, dtype, func):
+        a = testing.shaped_arange((5,), xp, numpy.int16)
+        return getattr(xp, func)(a, dtype=dtype)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_1d_include_initial(self, xp, dtype, func):
+        a = testing.shaped_arange((5,), xp, dtype)
+        return getattr(xp, func)(a, include_initial=True)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6, contiguous_check=False)
+    def test_axis_include_initial(self, xp, dtype, func):
+        a = testing.shaped_arange((3, 4, 5), xp, dtype)
+        return getattr(xp, func)(a, axis=1, include_initial=True)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_negative_axis_include_initial(self, xp, dtype, func):
+        a = testing.shaped_arange((3, 4, 5), xp, dtype)
+        return getattr(xp, func)(a, axis=-1, include_initial=True)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_out(self, xp, dtype, func):
+        a = testing.shaped_arange((3, 4, 5), xp, dtype)
+        out = xp.zeros((3, 4, 5), dtype=dtype)
+        getattr(xp, func)(a, axis=1, out=out)
+        return out
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_out_include_initial(self, xp, dtype, func):
+        a = testing.shaped_arange((3, 4, 5), xp, dtype)
+        out = xp.zeros((3, 5, 5), dtype=dtype)
+        getattr(xp, func)(a, axis=1, out=out, include_initial=True)
+        return out
+
+    @testing.for_all_dtypes()
+    def test_axis_none_requires_1d(self, dtype, func):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((3, 4), xp, dtype)
+            with pytest.raises(ValueError):
+                getattr(xp, func)(a)
+
+    @testing.for_all_dtypes()
+    def test_invalid_axis(self, dtype, func):
+        for xp in (numpy, cupy):
+            a = testing.shaped_arange((3, 4), xp, dtype)
+            with pytest.raises(AxisError):
+                getattr(xp, func)(a, axis=3)
+
+    def test_out_shape_mismatch(self, func):
+        a = testing.shaped_arange((3, 4), cupy, numpy.float32)
+        out = cupy.zeros((3, 5), dtype=numpy.float32)
+        with pytest.raises(ValueError):
+            getattr(cupy, func)(a, axis=1, out=out)
+
+    @pytest.mark.skip("implementation-defined acc to Python array API")
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_0d(self, xp, dtype, func):
+        # numpy accepts 0-D input; cupy matches via atleast_1d.
+        a = xp.asarray(3, dtype=dtype)
+        return getattr(xp, func)(a)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_0d_include_initial(self, xp, dtype, func):
+        a = xp.asarray(3, dtype=dtype)
+        return getattr(xp, func)(a, include_initial=True)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_empty_axis(self, xp, dtype, func):
+        a = xp.empty((3, 0, 4), dtype=dtype)
+        return getattr(xp, func)(a, axis=1)
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_empty_axis_include_initial(self, xp, dtype, func):
+        a = xp.empty((3, 0, 4), dtype=dtype)
+        return getattr(xp, func)(a, axis=1, include_initial=True)
+
+    @pytest.mark.parametrize("in_dtype", [numpy.int8, numpy.uint16])
+    @testing.numpy_cupy_array_equal()
+    def test_narrow_int_promotion_include_initial(self, xp, in_dtype, func):
+        # Narrow integer inputs must promote to int64 / uint64 (matches
+        # numpy's default-platform-integer rule) even when include_initial
+        # takes the internal-allocation path.
+        a = testing.shaped_arange((5,), xp, in_dtype)
+        return getattr(xp, func)(a, include_initial=True)
+
+    @testing.for_all_dtypes(no_bool=True)
+    @testing.numpy_cupy_allclose(rtol=1e-6)
+    def test_out_dtype_cast(self, xp, dtype, func):
+        # out has a different dtype than x -- result must be cast.
+        a = testing.shaped_arange((3, 4, 5), xp, numpy.int16)
+        out = xp.zeros((3, 4, 5), dtype=dtype)
+        getattr(xp, func)(a, axis=1, out=out)
+        return out