Add cumulative_prod (#795)

Author: tomwhite

.github/workflows/array-api-tests.yml

@@ -114,7 +114,8 @@ jobs:
           array_api_tests/test_array_object.py::test_getitem
           # test_searchsorted depends on sort which is not implemented
           array_api_tests/test_searching_functions.py::test_searchsorted
-          # cumulative_sum with include_initial=True is not implemented
+          # cumulative_* functions with include_initial=True are not implemented
+          array_api_tests/test_statistical_functions.py::test_cumulative_prod
           array_api_tests/test_statistical_functions.py::test_cumulative_sum
 
           # not implemented

api_status.md

@@ -79,7 +79,7 @@ This table shows which parts of the the [Array API](https://data-apis.org/array-
 |                          | `unique_values`     | :x:                |            | Shape is data dependent      |
 | Sorting Functions        | `argsort`           | :x:                |            |                              |
 |                          | `sort`              | :x:                |            |                              |
-| Statistical Functions    | `cumulative_prod`   | :x:                | 2024.12    |                              |
+| Statistical Functions    | `cumulative_prod`   | :white_check_mark: | 2024.12    |                              |
 |                          | `cumulative_sum`    | :white_check_mark: | 2023.12    |                              |
 |                          | `max`               | :white_check_mark: |            |                              |
 |                          | `mean`              | :white_check_mark: |            |                              |

cubed/__init__.py

@@ -341,6 +341,7 @@
 __all__ += ["argmax", "argmin", "count_nonzero", "searchsorted", "where"]
 
 from .array_api.statistical_functions import (
+    cumulative_prod,
     cumulative_sum,
     max,
     mean,
@@ -351,7 +352,17 @@
     var,
 )
 
-__all__ += ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
+__all__ += [
+    "cumulative_prod",
+    "cumulative_sum",
+    "max",
+    "mean",
+    "min",
+    "prod",
+    "std",
+    "sum",
+    "var",
+]
 
 from .array_api.utility_functions import all, any, diff
 

cubed/array_api/__init__.py

@@ -272,9 +272,29 @@
 
 __all__ += ["argmax", "argmin", "count_nonzero", "searchsorted", "where"]
 
-from .statistical_functions import cumulative_sum, max, mean, min, prod, std, sum, var
+from .statistical_functions import (
+    cumulative_prod,
+    cumulative_sum,
+    max,
+    mean,
+    min,
+    prod,
+    std,
+    sum,
+    var,
+)
 
-__all__ += ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
+__all__ += [
+    "cumulative_prod",
+    "cumulative_sum",
+    "max",
+    "mean",
+    "min",
+    "prod",
+    "std",
+    "sum",
+    "var",
+]
 
 from .utility_functions import all, any, diff
 

cubed/array_api/statistical_functions.py

@@ -10,6 +10,42 @@
 from cubed.core import reduction, scan
 
 
+def cumulative_prod(x, /, *, axis=None, dtype=None, include_initial=False, device=None):
+    if include_initial:
+        raise NotImplementedError("include_initial is not supported in cumulative_prod")
+    dtype = _upcast_integral_dtypes(
+        x,
+        dtype,
+        allowed_dtypes=(
+            "numeric",
+            "boolean",
+        ),
+        fname="cumulative_prod",
+        device=device,
+    )
+    return scan(
+        x,
+        preop=nxp.prod,
+        func=_cumulative_prod_func,
+        binop=nxp.multiply,
+        axis=axis,
+        dtype=dtype,
+    )
+
+
+def _cumulative_prod_func(a, /, *, axis=None, dtype=None, include_initial=False):
+    out = nxp.cumulative_prod(
+        a, axis=axis, dtype=dtype, include_initial=include_initial
+    )
+    if include_initial:
+        # we don't yet support including the final element as it complicates chunk sizing
+        ind = tuple(
+            slice(a.shape[i]) if i == axis else slice(None) for i in range(a.ndim)
+        )
+        out = out[ind]
+    return out
+
+
 def cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False, device=None):
     if include_initial:
         raise NotImplementedError("include_initial is not supported in cumulative_sum")

cubed/tests/test_array_api.py

@@ -850,6 +850,16 @@ def test_where_scalars():
 # Statistical functions
 
 
+@pytest.mark.parametrize("axis", [0, 1])
+def test_cumulative_prod_2d(axis):
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
+    b = xp.cumulative_prod(a, axis=axis)
+    assert_array_equal(
+        b.compute(),
+        np.cumulative_prod(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), axis=axis),
+    )
+
+
 @pytest.mark.parametrize("axis", [0, 1])
 def test_cumulative_sum_2d(axis):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))