scikit-learn
diff --git a/‎sklearn/metrics/_classification.py
Copy file name to clipboardExpand all lines: sklearn/metrics/_classification.py
+3-3Lines changed: 3 additions & 3 deletions b/‎sklearn/metrics/_classification.py
Copy file name to clipboardExpand all lines: sklearn/metrics/_classification.py
+3-3Lines changed: 3 additions & 3 deletions
diff --git a/‎sklearn/utils/_array_api.py
Copy file name to clipboardExpand all lines: sklearn/utils/_array_api.py
+11Lines changed: 11 additions & 0 deletions b/‎sklearn/utils/_array_api.py
Copy file name to clipboardExpand all lines: sklearn/utils/_array_api.py
+11Lines changed: 11 additions & 0 deletions
diff --git a/‎sklearn/utils/tests/test_array_api.py
Copy file name to clipboardExpand all lines: sklearn/utils/tests/test_array_api.py
+23-1Lines changed: 23 additions & 1 deletion b/‎sklearn/utils/tests/test_array_api.py
Copy file name to clipboardExpand all lines: sklearn/utils/tests/test_array_api.py
+23-1Lines changed: 23 additions & 1 deletion
@@ -28,6 +28,7 @@
 )
 from ..utils._array_api import (
     _average,
+    _count_nonzero,
     _is_numpy_namespace,
     _union1d,
     get_namespace,
@@ -221,9 +222,8 @@ def accuracy_score(y_true, y_pred, *, normalize=True, sample_weight=None):
         if _is_numpy_namespace(xp):
             differing_labels = count_nonzero(y_true - y_pred, axis=1)
         else:
-            differing_labels = xp.sum(
-                xp.astype(xp.astype(y_true - y_pred, xp.bool), xp.int8),
-                axis=1,
+            differing_labels = _count_nonzero(
+                y_true - y_pred, xp=xp, device=device, axis=1
             )
         score = xp.asarray(differing_labels == 0, device=device)
     else:
 
@@ -967,3 +967,14 @@ def _in1d(ar1, ar2, xp, assume_unique=False, invert=False):
         return ret[: ar1.shape[0]]
     else:
         return xp.take(ret, rev_idx, axis=0)
+
+
+def _count_nonzero(X, xp, device, axis=None):
+    """A variant of `sklearn.utils.sparsefuncs.count_nonzero` for the Array API."""
+    if axis == -1:
+        axis = 1
+    elif axis == -2:
+        axis = 0
+    one_scalar = xp.asarray(1, device=device)
+    zero_scalar = xp.asarray(0, device=device)
+    return xp.sum(xp.where(X != 0, one_scalar, zero_scalar), axis=axis)
@@ -13,6 +13,7 @@
     _atol_for_type,
     _average,
     _convert_to_numpy,
+    _count_nonzero,
     _estimator_with_converted_arrays,
     _is_numpy_namespace,
     _isin,
@@ -32,7 +33,7 @@
     assert_array_equal,
     skip_if_array_api_compat_not_configured,
 )
-from sklearn.utils.fixes import _IS_32BIT
+from sklearn.utils.fixes import _IS_32BIT, CSR_CONTAINERS
 
 
 @pytest.mark.parametrize("X", [numpy.asarray([1, 2, 3]), [1, 2, 3]])
@@ -566,3 +567,24 @@ def test_get_namespace_and_device():
         assert namespace is xp_torch
         assert is_array_api
         assert device == some_torch_tensor.device
+
+
+@pytest.mark.parametrize(
+    "array_namespace, device, _", yield_namespace_device_dtype_combinations()
+)
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+@pytest.mark.parametrize("axis", [0, 1, None, -1, -2])
+def test_count_nonzero(array_namespace, device, _, csr_container, axis):
+
+    from sklearn.utils.sparsefuncs import count_nonzero as sparse_count_nonzero
+
+    xp = _array_api_for_tests(array_namespace, device)
+    array = numpy.array([[0, 3, 0], [2, -1, 0], [0, 0, 0], [9, 8, 7], [4, 0, 5]])
+    expected = sparse_count_nonzero(csr_container(array), axis=axis)
+    array_xp = xp.asarray(array, device=device)
+
+    with config_context(array_api_dispatch=True):
+        result = _count_nonzero(array_xp, xp=xp, device=device, axis=axis)
+
+    assert_array_equal(_convert_to_numpy(result, xp=xp), expected)
+    assert getattr(array_xp, "device", None) == getattr(result, "device", None)