Adding diff function (#888)

Author: prady0t

ci/Numba-array-api-xfails.txt

@@ -69,13 +69,11 @@ array_api_tests/test_has_names.py::test_has_names[fft-irfftn]
 array_api_tests/test_creation_functions.py::test_empty_like
 array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 array_api_tests/test_manipulation_functions.py::test_squeeze
-array_api_tests/test_has_names.py::test_has_names[utility-diff]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_sum]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_prod]
 array_api_tests/test_has_names.py::test_has_names[indexing-take_along_axis]
 array_api_tests/test_has_names.py::test_has_names[searching-count_nonzero]
 array_api_tests/test_has_names.py::test_has_names[searching-searchsorted]
-array_api_tests/test_signatures.py::test_func_signature[diff]
 array_api_tests/test_signatures.py::test_func_signature[take_along_axis]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]

sparse/numba_backend/__init__.py

@@ -95,6 +95,7 @@
     can_cast,
     concat,
     concatenate,
+    diff,
     dot,
     einsum,
     empty,
@@ -341,6 +342,7 @@
     "repeat",
     "tile",
     "unstack",
+    "diff",
 ]
 
 

sparse/numba_backend/_common.py

@@ -3217,3 +3217,36 @@ def unstack(x, axis=0):
     new_order = (axis,) + tuple(i for i in range(ndim) if i != axis)
     x = x.transpose(new_order)
     return (*x,)
+
+
+def diff(x, axis=-1, n=1, prepend=None, append=None):
+    """
+    Calculates the n-th discrete difference along the given axis.
+
+    Parameters
+    ----------
+    x : SparseArray
+        Input sparse arrays.
+    n : int
+        The number of times values are differenced. Default: 1.
+    axis : int
+        The axis along which the difference is taken. Default: -1.
+
+    Returns
+    -------
+    out : SparseArray
+        An array containing the n-th discrete difference along the given axis.
+    """
+    if not isinstance(x, SparseArray):
+        raise TypeError("`x` must be a SparseArray.")
+
+    if axis < 0:
+        axis = x.ndim + axis
+    if prepend is not None:
+        x = concatenate([prepend, x], axis=axis)
+    if append is not None:
+        x = concatenate([x, append], axis=axis)
+    result = x
+    for _ in range(n):
+        result = result[(slice(None),) * axis + (slice(1, None),)] - result[(slice(None),) * axis + (slice(None, -1),)]
+    return result

sparse/numba_backend/tests/test_coo.py

@@ -2030,3 +2030,43 @@ def test_unstack_invalid_type():
     a = np.arange(6).reshape(2, 3)  # not a sparse array
     with pytest.raises(TypeError, match="must be a SparseArray"):
         sparse.unstack(a, axis=0)
+
+
+@pytest.mark.parametrize("ndim", range(1, 4))
+@pytest.mark.parametrize("shape_range", [3])
+@pytest.mark.parametrize("n", [1, 2])
+@pytest.mark.parametrize("use_prepend, use_append", [(False, False), (True, False), (False, True), (True, True)])
+def test_diff_matches_numpy(ndim, shape_range, n, use_prepend, use_append):
+    rng = np.random.default_rng(42)
+    shape = tuple(rng.integers(2, shape_range + 2) for _ in range(ndim))
+    x = rng.integers(0, 10, size=shape)
+    sparse_x = COO.from_numpy(x)
+
+    for axis in range(-ndim, ndim):
+        prepend = rng.integers(0, 10, size=x.shape).astype(x.dtype) if use_prepend else None
+        append = rng.integers(0, 10, size=x.shape).astype(x.dtype) if use_append else None
+
+        sparse_prepend = COO.from_numpy(prepend) if prepend is not None else None
+        sparse_append = COO.from_numpy(append) if append is not None else None
+
+        sparse_result = sparse.diff(sparse_x, axis=axis, n=n, prepend=sparse_prepend, append=sparse_append)
+
+        kwargs = {}
+        if prepend is not None:
+            kwargs["prepend"] = prepend
+        if append is not None:
+            kwargs["append"] = append
+
+        dense_result = np.diff(x, axis=axis, n=n, **kwargs)
+
+        np.testing.assert_array_equal(
+            sparse_result.todense(),
+            dense_result,
+            err_msg=f"Mismatch at axis={axis}, n={n}, prepend={use_prepend}, append={use_append}",
+        )
+
+
+def test_diff_invalid_type():
+    a = np.arange(6).reshape(2, 3)
+    with pytest.raises(TypeError, match="must be a SparseArray"):
+        sparse.diff(a)

sparse/numba_backend/tests/test_namespace.py

@@ -51,6 +51,7 @@ def test_namespace():
         "cosh",
         "diagonal",
         "diagonalize",
+        "diff",
         "divide",
         "dot",
         "e",