add a test for count_nonzero

Parrot the test from test_{argmin, argmax}

Author: ev-br

array_api_tests/test_searching_functions.py

@@ -88,6 +88,44 @@ def test_argmin(x, data):
         ph.assert_scalar_equals("argmin", type_=int, idx=out_idx, out=min_i, expected=expected)
 
 
+@pytest.mark.min_version("2024.12")
+@given(
+    x=hh.arrays(
+        dtype=hh.real_dtypes,
+        shape=hh.shapes(min_dims=1, min_side=1),
+        elements={"allow_nan": False},
+    ),
+    data=st.data(),
+)
+def test_count_nonzero(x, data):
+    kw = data.draw(
+        hh.kwargs(
+            axis=st.none() | st.integers(-x.ndim, max(x.ndim - 1, 0)),
+            keepdims=st.booleans(),
+        ),
+        label="kw",
+    )
+    keepdims = kw.get("keepdims", False)
+
+    out = xp.count_nonzero(x, **kw)
+
+    ph.assert_default_index("count_nonzero", out.dtype)
+    axes = sh.normalize_axis(kw.get("axis", None), x.ndim)
+    ph.assert_keepdimable_shape(
+        "count_nonzero", in_shape=x.shape, out_shape=out.shape, axes=axes, keepdims=keepdims, kw=kw
+    )
+    scalar_type = dh.get_scalar_type(x.dtype)
+
+    for indices, out_idx in zip(sh.axes_ndindex(x.shape, axes), sh.ndindex(out.shape)):
+        count = int(out[out_idx])
+        elements = []
+        for idx in indices:
+            s = scalar_type(x[idx])
+            elements.append(s)
+        expected = sum(el != 0 for el in elements)
+        ph.assert_scalar_equals("count_nonzero", type_=int, idx=out_idx, out=count, expected=expected)
+
+
 @given(hh.arrays(dtype=hh.all_dtypes, shape=()))
 def test_nonzero_zerodim_error(x):
     with pytest.raises(Exception):