Implements dpctl.tensor.count_nonzero

Author: ndgrigorian

dpctl/tensor/__init__.py

@@ -175,6 +175,7 @@
 from ._reduction import (
     argmax,
     argmin,
+    count_nonzero,
     logsumexp,
     max,
     min,
@@ -372,4 +373,5 @@
     "cumulative_prod",
     "cumulative_sum",
     "diff",
+    "count_nonzero",
 ]

dpctl/tensor/_reduction.py

@@ -772,3 +772,15 @@ def argmin(x, /, *, axis=None, keepdims=False, out=None):
             default array index data type for the device of ``x``.
     """
     return _search_over_axis(x, axis, keepdims, out, tri._argmin_over_axis)
+
+
+def count_nonzero(x, /, *, axis=None, keepdims=False, out=None):
+    if x.dtype != dpt.bool:
+        x = dpt.astype(x, dpt.bool, copy=False)
+    return sum(
+        x,
+        axis=axis,
+        dtype=ti.default_device_index_type(x.sycl_device),
+        keepdims=keepdims,
+        out=None,
+    )