feat: add cumulative_sum (#63)

Signed-off-by: nstarman <[email protected]>

src/quaxed/array_api/_statistical_functions.py

@@ -1,4 +1,4 @@
-__all__ = ["max", "mean", "min", "prod", "std", "sum", "var"]
+__all__ = ["cumulative_sum", "max", "mean", "min", "prod", "std", "sum", "var"]
 
 
 from jax.experimental import array_api
@@ -9,6 +9,20 @@
 from quaxed._utils import quaxify
 
 
+@quaxify
+def cumulative_sum(
+    x: ArrayLike,
+    /,
+    *,
+    axis: int | None = None,
+    dtype: DType | None = None,
+    include_initial: bool = False,
+) -> Value:
+    return array_api.cumulative_sum(
+        x, axis=axis, dtype=dtype, include_initial=include_initial
+    )
+
+
 @quaxify
 def max(  # pylint: disable=redefined-builtin
     x: ArrayLike,

tests/array_api/test_myarray.py

@@ -1180,6 +1180,42 @@ def test_sort():
 # Statistical functions
 
 
+def test_cumulative_sum():
+    """Test `cumulative_sum`."""
+    x = MyArray(xp.asarray([1, 2, 3], dtype=float))
+
+    # No arguments
+    got = xp.cumulative_sum(x)
+    expected = MyArray(xp.asarray([1, 3, 6], dtype=float))
+
+    assert isinstance(got, MyArray)
+    assert jnp.array_equal(got.array, expected.array)
+
+    # axis
+    got = xp.cumulative_sum(x, axis=0)
+    expected = MyArray(xp.asarray([1, 3, 6], dtype=float))
+
+    assert isinstance(got, MyArray)
+    assert jnp.array_equal(got.array, expected.array)
+
+    with pytest.raises(ValueError, match="axis 1"):
+        _ = xp.cumulative_sum(x, axis=1)
+
+    # dtype
+    got = xp.cumulative_sum(x, dtype=int)
+    expected = MyArray(xp.asarray([1, 3, 6], dtype=int))
+
+    assert isinstance(got, MyArray)
+    assert jnp.array_equal(got.array, expected.array)
+
+    # initial
+    got = xp.cumulative_sum(x, include_initial=True)
+    expected = MyArray(xp.asarray([0, 1, 3, 6]))
+
+    assert isinstance(got, MyArray)
+    assert jnp.array_equal(got.array, expected.array)
+
+
 @pytest.mark.skip("TODO")
 def test_max():
     """Test `max`."""

tests/myarray.py

@@ -264,16 +264,18 @@ def _complex_p(x: MyArray, y: MyArray) -> MyArray:
 def _concatenate_p(
     operand0: MyArray,
     *operands: MyArray,
-    dimension: Any,
+    **kwargs: Any,
 ) -> MyArray:
     return MyArray(
-        lax.concatenate(
-            [operand0.array] + [op.array for op in operands],
-            dimension=dimension,
-        ),
+        lax.concatenate([operand0.array] + [op.array for op in operands], **kwargs)
     )
 
 
+@register(lax.concatenate_p)
+def _concatenate_p(operand0: ArrayLike, operand1: MyArray, /, **kwargs: Any) -> MyArray:
+    return MyArray(lax.concatenate_p.bind(operand0, operand1.array, **kwargs))
+
+
 # ==============================================================================