Leverage on dpctl.tensor implementation in dpnp.count_nonzero (#1962)

* Leverage on dpctl.tensor implementation of count_nonzero

* Extend dpnp.get_result_array() to accept dpt.usm_ndarray

* Updated dpnp.mean() per review comment

* Add more dpnp tests to cover different use cases

dpnp/dpnp_iface.py

@@ -654,7 +654,7 @@ def get_result_array(a, out=None, casting="safe"):
 
     Parameters
     ----------
-    a : {dpnp_array}
+    a : {dpnp.ndarray, usm_ndarray}
         Input array.
     out : {dpnp.ndarray, usm_ndarray}
         If provided, value of `a` array will be copied into it
@@ -671,6 +671,8 @@ def get_result_array(a, out=None, casting="safe"):
     """
 
     if out is None:
+        if isinstance(a, dpt.usm_ndarray):
+            return dpnp_array._create_from_usm_ndarray(a)
         return a
 
     if isinstance(out, dpt.usm_ndarray):

dpnp/dpnp_iface_counting.py

@@ -44,25 +44,38 @@
 __all__ = ["count_nonzero"]
 
 
-def count_nonzero(a, axis=None, *, keepdims=False):
+def count_nonzero(a, axis=None, *, keepdims=False, out=None):
     """
     Counts the number of non-zero values in the array `a`.
 
     For full documentation refer to :obj:`numpy.count_nonzero`.
 
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        The array for which to count non-zeros.
+    axis : {None, int, tuple}, optional
+        Axis or tuple of axes along which to count non-zeros.
+        Default value means that non-zeros will be counted along a flattened
+        version of `a`.
+        Default: ``None``.
+    keepdims : bool, optional
+        If this is set to ``True``, the axes that are counted are left in the
+        result as dimensions with size one. With this option, the result will
+        broadcast correctly against the input array.
+        Default: ``False``.
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
+        The array into which the result is written. The data type of `out` must
+        match the expected shape and the expected data type of the result.
+        If ``None`` then a new array is returned.
+        Default: ``None``.
+
     Returns
     -------
     out : dpnp.ndarray
         Number of non-zero values in the array along a given axis.
-        Otherwise, a zero-dimensional array with the total number of
-        non-zero values in the array is returned.
-
-    Limitations
-    -----------
-    Parameters `a` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_ndarray`.
-    Otherwise ``TypeError`` exception will be raised.
-    Input array data types are limited by supported DPNP :ref:`Data types`.
+        Otherwise, a zero-dimensional array with the total number of non-zero
+        values in the array is returned.
 
     See Also
     --------
@@ -87,8 +100,10 @@ def count_nonzero(a, axis=None, *, keepdims=False):
 
     """
 
-    # TODO: might be improved by implementing an extension
-    # with `count_nonzero` kernel
     usm_a = dpnp.get_usm_ndarray(a)
-    usm_a = dpt.astype(usm_a, dpnp.bool, copy=False)
-    return dpnp.sum(usm_a, axis=axis, dtype=dpnp.intp, keepdims=keepdims)
+    usm_out = None if out is None else dpnp.get_usm_ndarray(out)
+
+    usm_res = dpt.count_nonzero(
+        usm_a, axis=axis, keepdims=keepdims, out=usm_out
+    )
+    return dpnp.get_result_array(usm_res, out)

dpnp/dpnp_iface_logic.py

@@ -51,7 +51,6 @@
 
 import dpnp
 from dpnp.dpnp_algo.dpnp_elementwise_common import DPNPBinaryFunc, DPNPUnaryFunc
-from dpnp.dpnp_array import dpnp_array
 
 __all__ = [
     "all",
@@ -167,13 +166,11 @@ def all(a, /, axis=None, out=None, keepdims=False, *, where=True):
 
     dpnp.check_limitations(where=where)
 
-    dpt_array = dpnp.get_usm_ndarray(a)
-    result = dpnp_array._create_from_usm_ndarray(
-        dpt.all(dpt_array, axis=axis, keepdims=keepdims)
-    )
+    usm_a = dpnp.get_usm_ndarray(a)
+    usm_res = dpt.all(usm_a, axis=axis, keepdims=keepdims)
+
     # TODO: temporary solution until dpt.all supports out parameter
-    result = dpnp.get_result_array(result, out)
-    return result
+    return dpnp.get_result_array(usm_res, out)
 
 
 def allclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False):
@@ -333,13 +330,11 @@ def any(a, /, axis=None, out=None, keepdims=False, *, where=True):
 
     dpnp.check_limitations(where=where)
 
-    dpt_array = dpnp.get_usm_ndarray(a)
-    result = dpnp_array._create_from_usm_ndarray(
-        dpt.any(dpt_array, axis=axis, keepdims=keepdims)
-    )
+    usm_a = dpnp.get_usm_ndarray(a)
+    usm_res = dpt.any(usm_a, axis=axis, keepdims=keepdims)
+
     # TODO: temporary solution until dpt.any supports out parameter
-    result = dpnp.get_result_array(result, out)
-    return result
+    return dpnp.get_result_array(usm_res, out)
 
 
 _EQUAL_DOCSTRING = """

dpnp/dpnp_iface_searching.py

@@ -391,7 +391,5 @@ def where(condition, x=None, y=None, /, *, order="K", out=None):
     usm_condition = dpnp.get_usm_ndarray(condition)
 
     usm_out = None if out is None else dpnp.get_usm_ndarray(out)
-    result = dpnp_array._create_from_usm_ndarray(
-        dpt.where(usm_condition, usm_x, usm_y, order=order, out=usm_out)
-    )
-    return dpnp.get_result_array(result, out)
+    usm_res = dpt.where(usm_condition, usm_x, usm_y, order=order, out=usm_out)
+    return dpnp.get_result_array(usm_res, out)

dpnp/dpnp_iface_statistics.py

@@ -614,13 +614,12 @@ def mean(a, /, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
 
     dpnp.check_limitations(where=where)
 
-    dpt_array = dpnp.get_usm_ndarray(a)
-    result = dpnp_array._create_from_usm_ndarray(
-        dpt.mean(dpt_array, axis=axis, keepdims=keepdims)
-    )
-    result = result.astype(dtype) if dtype is not None else result
+    usm_a = dpnp.get_usm_ndarray(a)
+    usm_res = dpt.mean(usm_a, axis=axis, keepdims=keepdims)
+    if dtype is not None:
+        usm_res = dpt.astype(usm_res, dtype)
 
-    return dpnp.get_result_array(result, out, casting="same_kind")
+    return dpnp.get_result_array(usm_res, out, casting="same_kind")
 
 
 def median(x1, axis=None, out=None, overwrite_input=False, keepdims=False):
@@ -904,11 +903,9 @@ def std(
         )
         dpnp.sqrt(result, out=result)
     else:
-        dpt_array = dpnp.get_usm_ndarray(a)
-        result = dpnp_array._create_from_usm_ndarray(
-            dpt.std(dpt_array, axis=axis, correction=ddof, keepdims=keepdims)
-        )
-        result = dpnp.get_result_array(result, out)
+        usm_a = dpnp.get_usm_ndarray(a)
+        usm_res = dpt.std(usm_a, axis=axis, correction=ddof, keepdims=keepdims)
+        result = dpnp.get_result_array(usm_res, out)
 
     if dtype is not None and out is None:
         result = result.astype(dtype, casting="same_kind")
@@ -1028,11 +1025,9 @@ def var(
 
         dpnp.divide(result, cnt, out=result)
     else:
-        dpt_array = dpnp.get_usm_ndarray(a)
-        result = dpnp_array._create_from_usm_ndarray(
-            dpt.var(dpt_array, axis=axis, correction=ddof, keepdims=keepdims)
-        )
-        result = dpnp.get_result_array(result, out)
+        usm_a = dpnp.get_usm_ndarray(a)
+        usm_res = dpt.var(usm_a, axis=axis, correction=ddof, keepdims=keepdims)
+        result = dpnp.get_result_array(usm_res, out)
 
     if out is None and dtype is not None:
         result = result.astype(dtype, casting="same_kind")

dpnp/dpnp_utils/dpnp_utils_reduction.py

@@ -25,7 +25,6 @@
 
 
 import dpnp
-from dpnp.dpnp_array import dpnp_array
 
 __all__ = ["dpnp_wrap_reduction_call"]
 
@@ -53,5 +52,4 @@ def dpnp_wrap_reduction_call(
 
     kwargs["out"] = usm_out
     res_usm = _reduction_fn(*args, **kwargs)
-    res = dpnp_array._create_from_usm_ndarray(res_usm)
-    return dpnp.get_result_array(res, input_out, casting="unsafe")
+    return dpnp.get_result_array(res_usm, input_out, casting="unsafe")

tests/test_counting.py

@@ -1,30 +1,124 @@
 import numpy
 import pytest
+from dpctl.tensor._numpy_helper import AxisError
 from numpy.testing import (
     assert_allclose,
+    assert_equal,
+    assert_raises,
 )
 
 import dpnp
 
 from .helper import (
     get_all_dtypes,
+    get_float_dtypes,
 )
 
 
-@pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
-@pytest.mark.parametrize("size", [2, 4, 8, 16, 3, 9, 27, 81])
-def test_count_nonzero(dtype, size):
-    if dtype != dpnp.bool:
-        a = numpy.arange(size, dtype=dtype)
-    else:
-        a = numpy.resize(numpy.arange(2, dtype=dtype), size)
+class TestCountNonZero:
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
+    @pytest.mark.parametrize("size", [2, 4, 8, 16, 3, 9, 27, 81])
+    def test_basic(self, dtype, size):
+        if dtype != dpnp.bool:
+            a = numpy.arange(size, dtype=dtype)
+        else:
+            a = numpy.resize(numpy.arange(2, dtype=dtype), size)
 
-    for i in range(int(size / 2)):
-        a[(i * (int(size / 3) - 1)) % size] = 0
+        for i in range(int(size / 2)):
+            a[(i * (int(size / 3) - 1)) % size] = 0
 
-    ia = dpnp.array(a)
+        ia = dpnp.array(a)
 
-    np_res = numpy.count_nonzero(a)
-    dpnp_res = dpnp.count_nonzero(ia)
+        result = dpnp.count_nonzero(ia)
+        expected = numpy.count_nonzero(a)
+        assert_allclose(result, expected)
 
-    assert_allclose(dpnp_res, np_res)
+    @pytest.mark.parametrize("data", [[], [0], [1]])
+    def test_trivial(self, data):
+        a = numpy.array(data)
+        ia = dpnp.array(a)
+
+        result = dpnp.count_nonzero(ia)
+        expected = numpy.count_nonzero(a)
+        assert_allclose(result, expected)
+
+    @pytest.mark.parametrize("data", [[], [0], [1]])
+    def test_trivial_boolean_dtype(self, data):
+        a = numpy.array(data, dtype="?")
+        ia = dpnp.array(a)
+
+        result = dpnp.count_nonzero(ia)
+        expected = numpy.count_nonzero(a)
+        assert_allclose(result, expected)
+
+    @pytest.mark.parametrize("axis", [0, 1])
+    def test_axis_basic(self, axis):
+        a = numpy.array([[0, 1, 7, 0, 0], [3, 0, 0, 2, 19]])
+        ia = dpnp.array(a)
+
+        result = dpnp.count_nonzero(ia, axis=axis)
+        expected = numpy.count_nonzero(a, axis=axis)
+        assert_equal(result, expected)
+
+    @pytest.mark.parametrize("xp", [numpy, dpnp])
+    def test_axis_raises(self, xp):
+        a = xp.array([[0, 1, 7, 0, 0], [3, 0, 0, 2, 19]])
+
+        assert_raises(ValueError, xp.count_nonzero, a, axis=(1, 1))
+        assert_raises(TypeError, xp.count_nonzero, a, axis="foo")
+        assert_raises(AxisError, xp.count_nonzero, a, axis=3)
+
+        # different exception type in numpy and dpnp
+        with pytest.raises((ValueError, TypeError)):
+            xp.count_nonzero(a, axis=xp.array([[1], [2]]))
+
+    @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True))
+    @pytest.mark.parametrize("axis", [0, 1, (0, 1), None])
+    def test_axis_all_dtypes(self, dt, axis):
+        a = numpy.zeros((3, 3), dtype=dt)
+        a[0, 0] = a[1, 0] = 1
+        ia = dpnp.array(a)
+
+        result = dpnp.count_nonzero(ia, axis=axis)
+        expected = numpy.count_nonzero(a, axis=axis)
+        assert_equal(result, expected)
+
+    def test_axis_empty(self):
+        axis = ()
+        a = numpy.array([[0, 0, 1], [1, 0, 1]])
+        ia = dpnp.array(a)
+
+        result = dpnp.count_nonzero(ia, axis=axis)
+        expected = numpy.count_nonzero(a, axis=axis)
+        assert_equal(result, expected)
+
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    def test_keepdims(self, axis):
+        a = numpy.array([[0, 0, 1, 0], [0, 3, 5, 0], [7, 9, 2, 0]])
+        ia = dpnp.array(a)
+
+        result = dpnp.count_nonzero(ia, axis=axis, keepdims=True)
+        expected = numpy.count_nonzero(a, axis=axis, keepdims=True)
+        assert_equal(result, expected)
+
+    @pytest.mark.parametrize("dt", get_all_dtypes(no_none=True))
+    def test_out(self, dt):
+        a = numpy.array([[0, 1, 0], [2, 0, 3]], dtype=dt)
+        ia = dpnp.array(a)
+        iout = dpnp.empty_like(ia, shape=ia.shape[1], dtype=dpnp.intp)
+
+        result = dpnp.count_nonzero(ia, axis=0, out=iout)
+        expected = numpy.count_nonzero(a, axis=0)  # no out keyword
+        assert_equal(result, expected)
+        assert result is iout
+
+    @pytest.mark.parametrize("dt", get_float_dtypes())
+    def test_out_floating_dtype(self, dt):
+        a = dpnp.array([[0, 1, 0], [2, 0, 3]])
+        out = dpnp.empty_like(a, shape=a.shape[1], dtype=dt)
+        assert_raises(ValueError, dpnp.count_nonzero, a, axis=0, out=out)
+
+    def test_array_method(self):
+        a = numpy.array([[1, 0, 0], [4, 0, 6]])
+        ia = dpnp.array(a)
+        assert_equal(ia.nonzero(), a.nonzero())