diff --git a/dpnp/dpnp_iface.py b/dpnp/dpnp_iface.py index bc005c21e0e..78e0342ef4c 100644 --- a/dpnp/dpnp_iface.py +++ b/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): diff --git a/dpnp/dpnp_iface_counting.py b/dpnp/dpnp_iface_counting.py index 515cad08a06..04de88c55c1 100644 --- a/dpnp/dpnp_iface_counting.py +++ b/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) diff --git a/dpnp/dpnp_iface_logic.py b/dpnp/dpnp_iface_logic.py index 2f28b3635ec..c6bc63f3ce7 100644 --- a/dpnp/dpnp_iface_logic.py +++ b/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 = """ diff --git a/dpnp/dpnp_iface_searching.py b/dpnp/dpnp_iface_searching.py index 88173b79ed7..455f862e614 100644 --- a/dpnp/dpnp_iface_searching.py +++ b/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) diff --git a/dpnp/dpnp_iface_statistics.py b/dpnp/dpnp_iface_statistics.py index ff12fb789b6..6e5ed6cc986 100644 --- a/dpnp/dpnp_iface_statistics.py +++ b/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") diff --git a/dpnp/dpnp_utils/dpnp_utils_reduction.py b/dpnp/dpnp_utils/dpnp_utils_reduction.py index 416f0fc2cdb..5565c8f1374 100644 --- a/dpnp/dpnp_utils/dpnp_utils_reduction.py +++ b/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") diff --git a/tests/test_counting.py b/tests/test_counting.py index c906fffcbe6..762abd58b68 100644 --- a/tests/test_counting.py +++ b/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())