Repository: pydata/sparse
Pull request #865 (2025-04-21 18:32:06)

[pre-commit.ci] pre-commit autoupdate
<!--pre-commit.ci start-->
updates:
- [github.com/astral-sh/ruff-pre-commit: v0.11.10 → v0.12.0](https://github.com/astral-sh/ruff-pre-commit/compare/v0.11.10...v0.12.0)
<!--pre-commit.ci end-->

hameerabbasi: cc @mtsokol @willow-ahrens Seems like we're ignoring `keepdims` here; assuming it's `True`. Was this related to the change discussed in the dev meeting two weeks ago?
codspeed-hq[bot]: <!-- __CODSPEED_PERFORMANCE_REPORT_COMMENT__ -->
<!-- __CODSPEED_INSTRUMENTATION_PERFORMANCE_REPORT_COMMENT__ -->

## [CodSpeed Performance Report](https://codspeed.io/pydata/sparse/branches/pre-commit-ci-update-config)

### Merging #865 will **improve performances by 14.37%**

<sub>Comparing <code>pre-commit-ci-update-config</code> (4f7de2f) with <code>main</code> (afb5212)</sub>



### Summary

`⚡ 10` improvements  
`✅ 330` untouched benchmarks  



### Benchmarks breakdown

|     | Benchmark | `BASE` | `HEAD` | Change |
| --- | --------- | ----------------------- | ------------------- | ------ |
| ⚡ | `` test_index_slice[side=100-rank=2-format='gcxs'] `` | 2.6 ms | 2.3 ms | +12.62% |
| ⚡ | `` test_elemwise[f=<built-in function add>-backend='SciPy'-side=1000] `` | 526.4 µs | 475.2 µs | +10.78% |
| ⚡ | `` test_elemwise[f=<built-in function add>-backend='SciPy'-side=100] `` | 462.4 µs | 411.6 µs | +12.32% |
| ⚡ | `` test_elemwise[f=<built-in function add>-backend='SciPy'-side=500] `` | 498.4 µs | 435.8 µs | +14.37% |
| ⚡ | `` test_elemwise[f=<built-in function gt>-backend='SciPy'-side=1000] `` | 524.7 µs | 473.3 µs | +10.85% |
| ⚡ | `` test_elemwise[f=<built-in function gt>-backend='SciPy'-side=100] `` | 462.4 µs | 410.8 µs | +12.58% |
| ⚡ | `` test_elemwise[f=<built-in function gt>-backend='SciPy'-side=500] `` | 478.3 µs | 426 µs | +12.27% |
| ⚡ | `` test_elemwise[f=<built-in function mul>-backend='SciPy'-side=1000] `` | 581.7 µs | 524.4 µs | +10.93% |
| ⚡ | `` test_elemwise[f=<built-in function mul>-backend='SciPy'-side=100] `` | 500.5 µs | 439.2 µs | +13.96% |
| ⚡ | `` test_elemwise[f=<built-in function mul>-backend='SciPy'-side=500] `` | 523.4 µs | 463.2 µs | +13.01% |

willow-ahrens: It looks like we might need to fix how keepdims is implemented, it looks like it works for eager but fails for lazy
willow-ahrens: By the way: Should these finch tests be moved to finch-tensor-python, or at least called from CI there?
willow-ahrens: I cannot reproduce the failing case locally. I am getting:

```
>>> A = finch.lazy(finch.Tensor(numpy.eye(5)))
>>> B = finch.lazy(finch.Tensor(sps.csr_matrix(numpy.eye(5))))
>>> finch.compute(finch.sum(finch.add(A, B), axis=0)).todense()
array([2., 2., 2., 2., 2.])
```
willow-ahrens: Is the problem just that we need to cut a new release of finch-tensor-python?
willow-ahrens: was that the fix then? Perhaps I should have bumped the breaking version number with the last finch.jl release. It was a bugfix but a breaking one. 
mtsokol: > was that the fix then? Perhaps I should have bumped the breaking version number with the last finch.jl release. It was a bugfix but a breaking one.

Yes, new release of `finch-tensor-python` fixed it! The previous version of `finch-tensor-python` was installed previously but newest `Finch.jl` so Python layer was missing the `keepdims` handling (I'm surprised that it installed newest `Finch.jl` as we also hardcode Finch version).

Numba Array API error is unrelated IMO, for Finch Array API I need to adjust skips and xfails file. 
mtsokol: Additional Numba and Finch Array API tests are due to new hypothesis version IMO. 
mtsokol: A few edge cases need to be resolved first to make Array API Finch job (described in https://github.com/finch-tensor/Finch.jl/issues/743)

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 124f520..85a1311 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -18,7 +18,7 @@ repos:
       exclude: ".ipynb"
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.11.10
+  rev: v0.12.0
   hooks:
     - id: ruff-check
       args: ["--fix"]
diff --git a/ci/Numba-array-api-xfails.txt b/ci/Numba-array-api-xfails.txt
index d745a73..66bd640 100644
--- a/ci/Numba-array-api-xfails.txt
+++ b/ci/Numba-array-api-xfails.txt
@@ -70,7 +70,6 @@ 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[manipulation-repeat]
 array_api_tests/test_has_names.py::test_has_names[manipulation-tile]
 array_api_tests/test_has_names.py::test_has_names[manipulation-unstack]
 array_api_tests/test_has_names.py::test_has_names[statistical-cumulative_sum]
@@ -79,7 +78,6 @@ 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[repeat]
 array_api_tests/test_signatures.py::test_func_signature[tile]
 array_api_tests/test_signatures.py::test_func_signature[unstack]
 array_api_tests/test_signatures.py::test_func_signature[take_along_axis]
@@ -107,7 +105,6 @@ array_api_tests/test_statistical_functions.py::test_cumulative_sum
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_1[None]
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[1]
 array_api_tests/test_array_object.py::test_getitem_arrays_and_ints_2[None]
-array_api_tests/test_manipulation_functions.py::test_repeat
 array_api_tests/test_searching_functions.py::test_count_nonzero
 array_api_tests/test_searching_functions.py::test_searchsorted
 array_api_tests/test_manipulation_functions.py::test_tile
diff --git a/sparse/numba_backend/__init__.py b/sparse/numba_backend/__init__.py
index 0702f5d..dae4b06 100644
--- a/sparse/numba_backend/__init__.py
+++ b/sparse/numba_backend/__init__.py
@@ -119,6 +119,7 @@ from ._common import (
     permute_dims,
     prod,
     real,
+    repeat,
     reshape,
     round,
     squeeze,
@@ -335,6 +336,7 @@ __all__ = [
     "where",
     "zeros",
     "zeros_like",
+    "repeat",
 ]
 
 
diff --git a/sparse/numba_backend/_common.py b/sparse/numba_backend/_common.py
index 23f07c6..abd581d 100644
--- a/sparse/numba_backend/_common.py
+++ b/sparse/numba_backend/_common.py
@@ -10,7 +10,7 @@ from numba import literal_unroll
 
 import numpy as np
 
-from ._coo import as_coo
+from ._coo import as_coo, expand_dims
 from ._sparse_array import SparseArray
 from ._utils import (
     _zero_of_dtype,
@@ -3104,3 +3104,46 @@ def vecdot(x1, x2, /, *, axis=-1):
         x1 = np.conjugate(x1)
 
     return np.sum(x1 * x2, axis=axis, dtype=np.result_type(x1, x2))
+
+
+def repeat(a, repeats, axis=None):
+    """
+    Repeat each element of an array after themselves
+
+    Parameters
+    ----------
+    a : SparseArray
+        Input sparse arrays
+    repeats : int
+        The number of repetitions for each element.
+        (Uneven repeats are not yet Implemented.)
+    axis : int, optional
+        The axis along which to repeat values. Returns a flattened sparse array if not specified.
+
+    Returns
+    -------
+    out : SparseArray
+        A sparse array which has the same shape as a, except along the given axis.
+    """
+    if not isinstance(a, SparseArray):
+        raise TypeError("`a` must be a SparseArray.")
+
+    if not isinstance(repeats, int):
+        raise ValueError("`repeats` must be an integer, uneven repeats are not yet Implemented.")
+    axes = list(range(a.ndim))
+    new_shape = list(a.shape)
+    axis_is_none = False
+    if axis is None:
+        a = a.reshape(-1)
+        axis = 0
+        axis_is_none = True
+    if axis < 0:
+        axis = a.ndim + axis
+    axes[a.ndim - 1], axes[axis] = axes[axis], axes[a.ndim - 1]
+    new_shape[axis] *= repeats
+    a = expand_dims(a, axis=axis + 1)
+    shape_to_broadcast = a.shape[: axis + 1] + (a.shape[axis + 1] * repeats,) + a.shape[axis + 2 :]
+    a = broadcast_to(a, shape_to_broadcast)
+    if not axis_is_none:
+        return a.reshape(new_shape)
+    return a.reshape(new_shape).flatten()


Test patch:
diff --git a/sparse/numba_backend/tests/test_coo.py b/sparse/numba_backend/tests/test_coo.py
index c31555f..ccc0bef 100644
--- a/sparse/numba_backend/tests/test_coo.py
+++ b/sparse/numba_backend/tests/test_coo.py
@@ -1926,3 +1926,33 @@ def test_xH_x():
     assert_eq(Ysp.conj().T @ Y, Y.conj().T @ Y)
     assert_eq(Ysp.conj().T @ Ysp, Y.conj().T @ Y)
     assert_eq(Y.conj().T @ Ysp.conj().T, Y.conj().T @ Y.conj().T)
+
+
+def test_repeat_invalid_input():
+    a = np.eye(3)
+    with pytest.raises(TypeError, match="`a` must be a SparseArray"):
+        sparse.repeat(a, repeats=2)
+    with pytest.raises(ValueError, match="`repeats` must be an integer"):
+        sparse.repeat(COO.from_numpy(a), repeats=[2, 2, 2])
+
+
+@pytest.mark.parametrize("ndim", range(1, 5))
+@pytest.mark.parametrize("repeats", [1, 2, 3])
+def test_repeat(ndim, repeats):
+    rng = np.random.default_rng()
+    shape = tuple(rng.integers(1, 4) for _ in range(ndim))
+    a = rng.integers(1, 10, size=shape)
+    sparse_a = COO.from_numpy(a)
+    for axis in [*range(-ndim, ndim), None]:
+        expected = np.repeat(a, repeats=repeats, axis=axis)
+        result_sparse = sparse.repeat(sparse_a, repeats=repeats, axis=axis)
+        actual = result_sparse.todense()
+        assert actual.shape == expected.shape, f"Shape mismatch on axis {axis}: {actual.shape} vs {expected.shape}"
+        np.testing.assert_array_equal(actual, expected)
+
+    expected = np.repeat(a, repeats=repeats, axis=None)
+    result_sparse = sparse.repeat(sparse_a, repeats=repeats, axis=None)
+    actual = result_sparse.todense()
+    print(f"Expected: {expected}, Actual: {actual}")
+    assert actual.shape == expected.shape
+    np.testing.assert_array_equal(actual, expected)
diff --git a/sparse/numba_backend/tests/test_namespace.py b/sparse/numba_backend/tests/test_namespace.py
index 4e86673..5aaa074 100644
--- a/sparse/numba_backend/tests/test_namespace.py
+++ b/sparse/numba_backend/tests/test_namespace.py
@@ -134,6 +134,7 @@ def test_namespace():
         "real",
         "reciprocal",
         "remainder",
+        "repeat",
         "reshape",
         "result_type",
         "roll",