ENH: add count_nonzero

Author: ev-br

array_api_strict/__init__.py

@@ -293,9 +293,9 @@
 
 __all__ += ["concat", "expand_dims", "flip", "moveaxis", "permute_dims", "repeat", "reshape", "roll", "squeeze", "stack", "tile", "unstack"]
 
-from ._searching_functions import argmax, argmin, nonzero, searchsorted, where
+from ._searching_functions import argmax, argmin, nonzero, count_nonzero, searchsorted, where
 
-__all__ += ["argmax", "argmin", "nonzero", "searchsorted", "where"]
+__all__ += ["argmax", "argmin", "nonzero", "count_nonzero", "searchsorted", "where"]
 
 from ._set_functions import unique_all, unique_counts, unique_inverse, unique_values
 

array_api_strict/_searching_functions.py

@@ -6,7 +6,7 @@
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from typing import Literal, Optional, Tuple
+    from typing import Literal, Optional, Tuple, Union
 
 import numpy as np
 
@@ -45,6 +45,24 @@ def nonzero(x: Array, /) -> Tuple[Array, ...]:
         raise ValueError("nonzero is not allowed on 0-dimensional arrays")
     return tuple(Array._new(i, device=x.device) for i in np.nonzero(x._array))
 
+
+@requires_api_version('2024.12')
+def count_nonzero(
+    x: Array,
+    /,
+    *,
+    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    keepdims: bool = False,
+) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.count_nonzero <numpy.count_nonzero>`
+
+    See its docstring for more information.
+    """
+    arr = np.count_nonzero(x._array, axis=axis, keepdims=keepdims)
+    return Array._new(np.asarray(arr), device=x.device)
+
+
 @requires_api_version('2023.12')
 def searchsorted(
     x1: Array,

array_api_strict/tests/test_flags.py

@@ -307,6 +307,7 @@ def test_api_version_2023_12(func_name):
     'reciprocal': lambda: xp.reciprocal(xp.asarray([2.])),
     'take_along_axis': lambda: xp.take_along_axis(xp.zeros((2, 3)),
                                                   xp.zeros((1, 4), dtype=xp.int64)),
+    'count_nonzero': lambda: xp.count_nonzero(xp.arange(3)),
 }
 
 @pytest.mark.parametrize('func_name', api_version_2024_12_examples.keys())