Add support of missing arguments in dpnp.count_nonzero (#1615)

* Added support of missing arguments in dpnp.count_nonzero

* Assert correct result dtype in test_count.py

* Addressed review comments

* Update test_count_nonzero to properly handle bool dtype

Author: antonwolfy

.github/workflows/conda-package.yml

@@ -45,6 +45,7 @@ env:
       third_party/cupy/math_tests/test_rounding.py
       third_party/cupy/math_tests/test_trigonometric.py
       third_party/cupy/sorting_tests/test_sort.py
+      third_party/cupy/sorting_tests/test_count.py
   VER_JSON_NAME: 'version.json'
   VER_SCRIPT1: "import json; f = open('version.json', 'r'); j = json.load(f); f.close(); "
   VER_SCRIPT2: "d = j['dpnp'][0]; print('='.join((d[s] for s in ('version', 'build'))))"

dpnp/backend/cmake/Modules/IntelSYCLConfig.cmake

@@ -18,7 +18,7 @@
 IntelSYCLConfig
 -------
 
-Library to verify SYCL compatability of CMAKE_CXX_COMPILER
+Library to verify SYCL compatibility of CMAKE_CXX_COMPILER
 and passes relevant compiler flags.
 
 Result Variables
@@ -290,7 +290,7 @@ if(nosycllang)
   set(IntelSYCL_NOT_FOUND_MESSAGE "${SYCL_REASON_FAILURE}")
 endif()
 
-# Placeholder for identifying various implemenations of SYCL compilers.
+# Placeholder for identifying various implementations of SYCL compilers.
 # for now, set to the CMAKE_CXX_COMPILER_ID
 set(SYCL_IMPLEMENTATION_ID "${CMAKE_CXX_COMPILER_ID}")
 

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -115,25 +115,23 @@ enum class DPNPFuncName : size_t
     DPNP_FN_COS,           /**< Used in numpy.cos() impl  */
     DPNP_FN_COSH,          /**< Used in numpy.cosh() impl  */
     DPNP_FN_COUNT_NONZERO, /**< Used in numpy.count_nonzero() impl  */
-    DPNP_FN_COUNT_NONZERO_EXT, /**< Used in numpy.count_nonzero() impl, requires
-                                  extra parameters */
-    DPNP_FN_COV,               /**< Used in numpy.cov() impl  */
-    DPNP_FN_CROSS,             /**< Used in numpy.cross() impl  */
-    DPNP_FN_CROSS_EXT,         /**< Used in numpy.cross() impl, requires extra
-                                  parameters */
-    DPNP_FN_CUMPROD,           /**< Used in numpy.cumprod() impl  */
-    DPNP_FN_CUMPROD_EXT,       /**< Used in numpy.cumprod() impl, requires extra
-                                  parameters */
-    DPNP_FN_CUMSUM,            /**< Used in numpy.cumsum() impl  */
-    DPNP_FN_CUMSUM_EXT,        /**< Used in numpy.cumsum() impl, requires extra
-                                  parameters */
-    DPNP_FN_DEGREES,           /**< Used in numpy.degrees() impl  */
-    DPNP_FN_DEGREES_EXT,       /**< Used in numpy.degrees() impl, requires extra
-                                  parameters */
-    DPNP_FN_DET,               /**< Used in numpy.linalg.det() impl  */
-    DPNP_FN_DET_EXT,  /**< Used in numpy.linalg.det() impl, requires extra
-                         parameters */
-    DPNP_FN_DIAG,     /**< Used in numpy.diag() impl  */
+    DPNP_FN_COV,           /**< Used in numpy.cov() impl  */
+    DPNP_FN_CROSS,         /**< Used in numpy.cross() impl  */
+    DPNP_FN_CROSS_EXT,     /**< Used in numpy.cross() impl, requires extra
+                              parameters */
+    DPNP_FN_CUMPROD,       /**< Used in numpy.cumprod() impl  */
+    DPNP_FN_CUMPROD_EXT,   /**< Used in numpy.cumprod() impl, requires extra
+                              parameters */
+    DPNP_FN_CUMSUM,        /**< Used in numpy.cumsum() impl  */
+    DPNP_FN_CUMSUM_EXT,    /**< Used in numpy.cumsum() impl, requires extra
+                              parameters */
+    DPNP_FN_DEGREES,       /**< Used in numpy.degrees() impl  */
+    DPNP_FN_DEGREES_EXT,   /**< Used in numpy.degrees() impl, requires extra
+                              parameters */
+    DPNP_FN_DET,           /**< Used in numpy.linalg.det() impl  */
+    DPNP_FN_DET_EXT,       /**< Used in numpy.linalg.det() impl, requires extra
+                              parameters */
+    DPNP_FN_DIAG,          /**< Used in numpy.diag() impl  */
     DPNP_FN_DIAG_EXT, /**< Used in numpy.diag() impl, requires extra parameters
                        */
     DPNP_FN_DIAG_INDICES,     /**< Used in numpy.diag_indices() impl  */

dpnp/backend/kernels/dpnp_krnl_statistics.cpp

@@ -292,14 +292,6 @@ template <typename _DataType_input, typename _DataType_output>
 void (*dpnp_count_nonzero_default_c)(void *, void *, size_t) =
     dpnp_count_nonzero_c<_DataType_input, _DataType_output>;
 
-template <typename _DataType_input, typename _DataType_output>
-DPCTLSyclEventRef (*dpnp_count_nonzero_ext_c)(DPCTLSyclQueueRef,
-                                              void *,
-                                              void *,
-                                              size_t,
-                                              const DPCTLEventVectorRef) =
-    dpnp_count_nonzero_c<_DataType_input, _DataType_output>;
-
 template <typename _DataType>
 class dpnp_max_c_kernel;
 
@@ -1273,17 +1265,6 @@ void func_map_init_statistics(func_map_t &fmap)
     fmap[DPNPFuncName::DPNP_FN_COUNT_NONZERO][eft_DBL][eft_DBL] = {
         eft_LNG, (void *)dpnp_count_nonzero_default_c<double, int64_t>};
 
-    fmap[DPNPFuncName::DPNP_FN_COUNT_NONZERO_EXT][eft_BLN][eft_BLN] = {
-        eft_LNG, (void *)dpnp_count_nonzero_ext_c<bool, int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_COUNT_NONZERO_EXT][eft_INT][eft_INT] = {
-        eft_LNG, (void *)dpnp_count_nonzero_ext_c<int32_t, int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_COUNT_NONZERO_EXT][eft_LNG][eft_LNG] = {
-        eft_LNG, (void *)dpnp_count_nonzero_ext_c<int64_t, int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_COUNT_NONZERO_EXT][eft_FLT][eft_FLT] = {
-        eft_LNG, (void *)dpnp_count_nonzero_ext_c<float, int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_COUNT_NONZERO_EXT][eft_DBL][eft_DBL] = {
-        eft_LNG, (void *)dpnp_count_nonzero_ext_c<double, int64_t>};
-
     fmap[DPNPFuncName::DPNP_FN_COV][eft_INT][eft_INT] = {
         eft_DBL, (void *)dpnp_cov_default_c<double>};
     fmap[DPNPFuncName::DPNP_FN_COV][eft_LNG][eft_LNG] = {

dpnp/dpnp_algo/CMakeLists.txt

@@ -2,7 +2,6 @@
 set(dpnp_algo_pyx_deps
   ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_linearalgebra.pxi
   ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_manipulation.pxi
-  ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_counting.pxi
   ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_statistics.pxi
   ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_trigonometric.pxi
   ${CMAKE_CURRENT_SOURCE_DIR}/dpnp_algo_sorting.pxi

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -54,8 +54,6 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_COPYSIGN_EXT
         DPNP_FN_CORRELATE
         DPNP_FN_CORRELATE_EXT
-        DPNP_FN_COUNT_NONZERO
-        DPNP_FN_COUNT_NONZERO_EXT
         DPNP_FN_CROSS
         DPNP_FN_CROSS_EXT
         DPNP_FN_CUMPROD

dpnp/dpnp_algo/dpnp_algo.pyx

@@ -60,7 +60,6 @@ __all__ = [
 
 
 include "dpnp_algo_arraycreation.pxi"
-include "dpnp_algo_counting.pxi"
 include "dpnp_algo_indexing.pxi"
 include "dpnp_algo_linearalgebra.pxi"
 include "dpnp_algo_logic.pxi"

dpnp/dpnp_algo/dpnp_algo_counting.pxi

@@ -39,48 +39,54 @@
 
 """
 
-
-import numpy
-
 import dpnp
-from dpnp.dpnp_algo.dpnp_algo import *  # TODO need to investigate why dpnp.dpnp_algo can not be used
-from dpnp.dpnp_utils import *
 
 __all__ = ["count_nonzero"]
 
 
-def count_nonzero(x1, axis=None, *, keepdims=False):
+def count_nonzero(a, axis=None, *, keepdims=False):
     """
-    Counts the number of non-zero values in the array ``in_array1``.
+    Counts the number of non-zero values in the array `a`.
 
     For full documentation refer to :obj:`numpy.count_nonzero`.
 
+    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
     -----------
-    Parameter `x1` is supported as :obj:`dpnp.ndarray`.
-    Otherwise the function will be executed sequentially on CPU.
-    Parameter `axis` is supported only with default value `None`.
-    Parameter `keepdims` is supported only with default value `False`.
+    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`.
+
+    See Also
+    --------
+    :obj:`dpnp.nonzero` : Return the coordinates of all the non-zero values.
 
     Examples
     --------
     >>> import dpnp as np
-    >>> np.count_nonzero(np.array([1, 0, 3, 0, 5])
-    3
-    >>> np.count_nonzero(np.array([[1, 0, 3, 0, 5],[0, 9, 0, 7, 0]]))
-    5
+    >>> np.count_nonzero(np.eye(4))
+    array(4)
+    >>> a = np.array([[0, 1, 7, 0],
+                      [3, 0, 2, 19]])
+    >>> np.count_nonzero(a)
+    array(5)
+    >>> np.count_nonzero(a, axis=0)
+    array([1, 1, 2, 1])
+    >>> np.count_nonzero(a, axis=1)
+    array([2, 3])
+    >>> np.count_nonzero(a, axis=1, keepdims=True)
+    array([[2],
+           [3]])
 
     """
-    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
-    if x1_desc:
-        if axis is not None:
-            pass
-        elif keepdims is not False:
-            pass
-        else:
-            result_obj = dpnp_count_nonzero(x1_desc).get_pyobj()
-            result = dpnp.convert_single_elem_array_to_scalar(result_obj)
-
-            return result
 
-    return call_origin(numpy.count_nonzero, x1, axis, keepdims=keepdims)
+    # TODO: might be improved by implementing an extension with `count_nonzero` kernel
+    a = dpnp.astype(a, dpnp.bool, copy=False)
+    return a.sum(axis=axis, dtype=dpnp.intp, keepdims=keepdims)

dpnp/dpnp_iface_counting.py

@@ -72,6 +72,7 @@
     "int64",
     "integer",
     "intc",
+    "intp",
     "isscalar",
     "issubdtype",
     "issubsctype",
@@ -119,6 +120,7 @@
 int64 = numpy.int64
 integer = numpy.integer
 intc = numpy.intc
+intp = numpy.intp
 number = numpy.number
 signedinteger = numpy.signedinteger
 single = numpy.single

dpnp/dpnp_iface_types.py

@@ -815,10 +815,6 @@ tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers::test_hig
 tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers::test_normal
 tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers::test_size_is_not_none
 
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero_int_axis
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero_tuple_axis
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero_zero_dim
 tests/third_party/cupy/sorting_tests/test_search.py::TestArgMinMaxDtype_param_0_{func='argmin', is_module=True, shape=(3, 4)}::test_argminmax_dtype
 tests/third_party/cupy/sorting_tests/test_search.py::TestArgMinMaxDtype_param_1_{func='argmin', is_module=True, shape=()}::test_argminmax_dtype
 tests/third_party/cupy/sorting_tests/test_search.py::TestArgMinMaxDtype_param_2_{func='argmin', is_module=False, shape=(3, 4)}::test_argminmax_dtype

tests/skipped_tests.tbl

@@ -902,10 +902,6 @@ tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_bo
 tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_goodness_of_fit
 tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_goodness_of_fit_2
 
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero_int_axis
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero_tuple_axis
-tests/third_party/cupy/sorting_tests/test_count.py::TestCount::test_count_nonzero_zero_dim
 tests/third_party/cupy/sorting_tests/test_search.py::TestArgMinMaxDtype_param_0_{func='argmin', is_module=True, shape=(3, 4)}::test_argminmax_dtype
 tests/third_party/cupy/sorting_tests/test_search.py::TestArgMinMaxDtype_param_1_{func='argmin', is_module=True, shape=()}::test_argminmax_dtype
 tests/third_party/cupy/sorting_tests/test_search.py::TestArgMinMaxDtype_param_2_{func='argmin', is_module=False, shape=(3, 4)}::test_argminmax_dtype

tests/skipped_tests_gpu.tbl

@@ -1,17 +1,24 @@
 import numpy
 import pytest
+from numpy.testing import (
+    assert_allclose,
+)
 
 import dpnp
 
-
-@pytest.mark.parametrize(
-    "type",
-    [numpy.float64, numpy.float32, numpy.int64, numpy.int32],
-    ids=["float64", "float32", "int64", "int32"],
+from .helper import (
+    get_all_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(type, size):
-    a = numpy.arange(size, dtype=type)
+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)
+
     for i in range(int(size / 2)):
         a[(i * (int(size / 3) - 1)) % size] = 0
 
@@ -20,4 +27,4 @@ def test_count_nonzero(type, size):
     np_res = numpy.count_nonzero(a)
     dpnp_res = dpnp.count_nonzero(ia)
 
-    numpy.testing.assert_allclose(dpnp_res, np_res)
+    assert_allclose(dpnp_res, np_res)

tests/test_counting.py

@@ -238,11 +238,12 @@ def test_meshgrid(device_x, device_y):
         pytest.param("arctanh", [-0.5, 0.0, 0.5]),
         pytest.param("ceil", [-1.7, -1.5, -0.2, 0.2, 1.5, 1.7, 2.0]),
         pytest.param("conjugate", [[1.0 + 1.0j, 0.0], [0.0, 1.0 + 1.0j]]),
+        pytest.param("copy", [1.0, 2.0, 3.0]),
         pytest.param(
             "cos", [-dpnp.pi / 2, -dpnp.pi / 4, 0.0, dpnp.pi / 4, dpnp.pi / 2]
         ),
         pytest.param("cosh", [-5.0, -3.5, 0.0, 3.5, 5.0]),
-        pytest.param("copy", [1.0, 2.0, 3.0]),
+        pytest.param("count_nonzero", [3, 0, 2, -1.2]),
         pytest.param("cumprod", [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
         pytest.param("cumsum", [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),
         pytest.param("diff", [1.0, 2.0, 4.0, 7.0, 0.0]),

tests/test_sycl_queue.py

@@ -338,6 +338,7 @@ def test_meshgrid(usm_type_x, usm_type_y):
             "cos", [-dp.pi / 2, -dp.pi / 4, 0.0, dp.pi / 4, dp.pi / 2]
         ),
         pytest.param("cosh", [-5.0, -3.5, 0.0, 3.5, 5.0]),
+        pytest.param("count_nonzero", [0, 1, 7, 0]),
         pytest.param("exp", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("expm1", [1.0e-10, 1.0, 2.0, 4.0, 7.0]),
         pytest.param("floor", [-1.7, -1.5, -0.2, 0.2, 1.5, 1.7, 2.0]),

tests/test_usm_type.py

@@ -6,7 +6,6 @@
 from tests.third_party.cupy import testing
 
 
-@testing.gpu
 class TestCount(unittest.TestCase):
     @testing.for_all_dtypes()
     def test_count_nonzero(self, dtype):
@@ -17,12 +16,12 @@ def func(xp):
             if xp is cupy:
                 # CuPy returns zero-dimensional array instead of
                 # returning a scalar value
-                self.assertIsInstance(c, xp.ndarray)
-                self.assertEqual(c.dtype, "l")
-                self.assertEqual(c.shape, ())
+                assert isinstance(c, xp.ndarray)
+                assert c.dtype == "p"
+                assert c.shape == ()
             return int(c)
 
-        self.assertEqual(func(numpy), func(cupy))
+        assert func(numpy) == func(cupy)
 
     @testing.for_all_dtypes()
     def test_count_nonzero_zero_dim(self, dtype):
@@ -32,12 +31,12 @@ def func(xp):
             if xp is cupy:
                 # CuPy returns zero-dimensional array instead of
                 # returning a scalar value
-                self.assertIsInstance(c, xp.ndarray)
-                self.assertEqual(c.dtype, "l")
-                self.assertEqual(c.shape, ())
+                assert isinstance(c, xp.ndarray)
+                assert c.dtype == "p"
+                assert c.shape == ()
             return int(c)
 
-        self.assertEqual(func(numpy), func(cupy))
+        assert func(numpy) == func(cupy)
 
     @testing.for_all_dtypes()
     def test_count_nonzero_int_axis(self, dtype):