implement dpnp.mean (#1632)

* implement dpnp.mean

* address comments

Author: vtavana

.github/workflows/conda-package.yml

@@ -25,6 +25,7 @@ env:
       test_random_state.py
       test_sort.py
       test_special.py
+      test_statistics.py
       test_sycl_queue.py
       test_umath.py
       test_usm_type.py
@@ -47,6 +48,7 @@ env:
       third_party/cupy/math_tests/test_trigonometric.py
       third_party/cupy/sorting_tests/test_sort.py
       third_party/cupy/sorting_tests/test_count.py
+      third_party/cupy/statistics_tests/test_meanvar.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/dpnp_array.py

@@ -915,10 +915,16 @@ def max(
 
         return dpnp.max(self, axis, out, keepdims, initial, where)
 
-    def mean(self, axis=None, **kwargs):
-        """Returns the average of the array elements."""
+    def mean(
+        self, axis=None, dtype=None, out=None, keepdims=False, *, where=True
+    ):
+        """
+        Returns the average of the array elements.
+
+        Refer to :obj:`dpnp.mean` for full documentation.
+        """
 
-        return dpnp.mean(self, axis=axis, **kwargs)
+        return dpnp.mean(self, axis, dtype, out, keepdims, where=where)
 
     def min(
         self,

dpnp/dpnp_iface_statistics.py

@@ -40,7 +40,6 @@
 
 import dpctl.tensor as dpt
 import numpy
-from numpy.core.numeric import normalize_axis_tuple
 
 import dpnp
 from dpnp.dpnp_algo import *
@@ -417,24 +416,24 @@ def max(a, axis=None, out=None, keepdims=False, initial=None, where=True):
         return dpnp.get_result_array(result, out)
 
 
-def mean(x, /, *, axis=None, dtype=None, keepdims=False, out=None, where=True):
+def mean(a, /, axis=None, dtype=None, out=None, keepdims=False, *, where=True):
     """
     Compute the arithmetic mean along the specified axis.
 
     For full documentation refer to :obj:`numpy.mean`.
 
     Returns
     -------
-    y : dpnp.ndarray
+    out : dpnp.ndarray
         an array containing the mean values of the elements along the specified axis(axes).
-        If the input array is empty, an array containing a single NaN value is returned.
+        If the input is a zero-size array, an array containing NaN values is returned.
 
     Limitations
     -----------
-    Parameters `x` is supported as either :class:`dpnp.ndarray`
+    Parameters `a` is supported as either :class:`dpnp.ndarray`
     or :class:`dpctl.tensor.usm_ndarray`.
-    Parameters `keepdims`, `out` and `where` are supported with their default values.
-    Otherwise the function will be executed sequentially on CPU.
+    Parameter `where` is supported only with their default values.
+    Otherwise ``NotImplementedError`` exception will be raised.
     Input array data types are limited by supported DPNP :ref:`Data types`.
 
     See Also
@@ -459,59 +458,21 @@ def mean(x, /, *, axis=None, dtype=None, keepdims=False, out=None, where=True):
     array([2., 3.])
     >>> np.mean(a, axis=1)
     array([1.5, 3.5])
+
     """
 
-    if keepdims is not False:
-        pass
-    elif out is not None:
-        pass
-    elif where is not True:
-        pass
+    if where is not True:
+        raise NotImplementedError(
+            "where keyword argument is only supported by its default value."
+        )
     else:
-        if dtype is None and dpnp.issubdtype(x.dtype, dpnp.inexact):
-            dtype = x.dtype
-
-        if axis is None:
-            if x.size == 0:
-                return dpnp.array(dpnp.nan, dtype=dtype)
-            else:
-                result = dpnp.sum(x, dtype=dtype) / x.size
-                return result.astype(dtype) if result.dtype != dtype else result
-
-        if not isinstance(axis, (tuple, list)):
-            axis = (axis,)
-
-        axis = normalize_axis_tuple(axis, x.ndim, "axis")
-        res_sum = dpnp.sum(x, axis=axis, dtype=dtype)
-
-        del_ = 1.0
-        for axis_value in axis:
-            del_ *= x.shape[axis_value]
-
-        # performing an inplace operation on arrays of bool or integer types
-        # is not possible due to incompatible data types because
-        # it returns a floating value
-        if dpnp.issubdtype(res_sum.dtype, dpnp.inexact):
-            res_sum /= del_
-        else:
-            new_res_sum = res_sum / del_
-            return (
-                new_res_sum.astype(dtype)
-                if new_res_sum.dtype != dtype
-                else new_res_sum
-            )
-
-        return res_sum.astype(dtype) if res_sum.dtype != dtype else res_sum
+        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
 
-    return call_origin(
-        numpy.mean,
-        x,
-        axis=axis,
-        dtype=dtype,
-        out=out,
-        keepdims=keepdims,
-        where=where,
-    )
+        return dpnp.get_result_array(result, out)
 
 
 def median(x1, axis=None, out=None, overwrite_input=False, keepdims=False):

tests/test_mathematical.py

@@ -1993,45 +1993,6 @@ def test_sum(shape, dtype_in, dtype_out, transpose, keepdims, order):
         assert_array_equal(numpy_res, dpnp_res.asnumpy())
 
 
-class TestMean:
-    @pytest.mark.parametrize("dtype", get_all_dtypes())
-    def test_mean_axis_tuple(self, dtype):
-        dp_array = dpnp.array([[0, 1, 2], [3, 4, 0]], dtype=dtype)
-        np_array = dpnp.asnumpy(dp_array)
-
-        result = dpnp.mean(dp_array, axis=(0, 1))
-        expected = numpy.mean(np_array, axis=(0, 1))
-        assert_allclose(expected, result)
-
-    def test_mean_axis_zero_size(self):
-        dp_array = dpnp.array([], dtype="int64")
-        np_array = dpnp.asnumpy(dp_array)
-
-        result = dpnp.mean(dp_array)
-        expected = numpy.mean(np_array)
-        assert_allclose(expected, result)
-
-    def test_mean_strided(self):
-        dp_array = dpnp.array([-2, -1, 0, 1, 0, 2], dtype="f4")
-        np_array = dpnp.asnumpy(dp_array)
-
-        result = dpnp.mean(dp_array[::-1])
-        expected = numpy.mean(np_array[::-1])
-        assert_allclose(expected, result)
-
-        result = dpnp.mean(dp_array[::2])
-        expected = numpy.mean(np_array[::2])
-        assert_allclose(expected, result)
-
-    def test_mean_scalar(self):
-        dp_array = dpnp.array(5)
-        np_array = dpnp.asnumpy(dp_array)
-
-        result = dp_array.mean()
-        expected = np_array.mean()
-        assert_allclose(expected, result)
-
-
 @pytest.mark.parametrize(
     "dtype", get_all_dtypes(no_bool=True, no_none=True, no_complex=True)
 )

tests/test_statistics.py

@@ -8,7 +8,7 @@
 
 import dpnp
 
-from .helper import get_all_dtypes
+from .helper import assert_dtype_allclose, get_all_dtypes
 
 
 @pytest.mark.parametrize(
@@ -88,6 +88,73 @@ def test_max_min_NotImplemented(func):
         getattr(dpnp, func)(ia, initial=6)
 
 
+class TestMean:
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    def test_mean_axis_tuple(self, dtype):
+        dp_array = dpnp.array([[0, 1, 2], [3, 4, 0]], dtype=dtype)
+        np_array = dpnp.asnumpy(dp_array)
+
+        result = dpnp.mean(dp_array, axis=(0, 1))
+        expected = numpy.mean(np_array, axis=(0, 1))
+        assert_allclose(expected, result)
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    @pytest.mark.parametrize("axis", [0, 1, (0, 1)])
+    def test_mean_out(self, dtype, axis):
+        dp_array = dpnp.array([[0, 1, 2], [3, 4, 0]], dtype=dtype)
+        np_array = dpnp.asnumpy(dp_array)
+
+        expected = numpy.mean(np_array, axis=axis)
+        result = dpnp.empty_like(dpnp.asarray(expected))
+        dpnp.mean(dp_array, axis=axis, out=result)
+        assert_dtype_allclose(result, expected)
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    def test_mean_dtype(self, dtype):
+        dp_array = dpnp.array([[0, 1, 2], [3, 4, 0]], dtype="i4")
+        np_array = dpnp.asnumpy(dp_array)
+
+        expected = numpy.mean(np_array, dtype=dtype)
+        result = dpnp.mean(dp_array, dtype=dtype)
+        assert_allclose(expected, result)
+
+    @pytest.mark.usefixtures("suppress_invalid_numpy_warnings")
+    @pytest.mark.parametrize("axis", [0, 1, (0, 1)])
+    @pytest.mark.parametrize("shape", [(2, 3), (2, 0), (0, 3)])
+    def test_mean_empty(self, axis, shape):
+        dp_array = dpnp.empty(shape, dtype=dpnp.int64)
+        np_array = dpnp.asnumpy(dp_array)
+
+        result = dpnp.mean(dp_array, axis=axis)
+        expected = numpy.mean(np_array, axis=axis)
+        assert_allclose(expected, result)
+
+    def test_mean_strided(self):
+        dp_array = dpnp.array([-2, -1, 0, 1, 0, 2], dtype="f4")
+        np_array = dpnp.asnumpy(dp_array)
+
+        result = dpnp.mean(dp_array[::-1])
+        expected = numpy.mean(np_array[::-1])
+        assert_allclose(expected, result)
+
+        result = dpnp.mean(dp_array[::2])
+        expected = numpy.mean(np_array[::2])
+        assert_allclose(expected, result)
+
+    def test_mean_scalar(self):
+        dp_array = dpnp.array(5)
+        np_array = dpnp.asnumpy(dp_array)
+
+        result = dp_array.mean()
+        expected = np_array.mean()
+        assert_allclose(expected, result)
+
+    def test_mean_NotImplemented(func):
+        ia = dpnp.arange(5)
+        with pytest.raises(NotImplementedError):
+            dpnp.mean(ia, where=False)
+
+
 @pytest.mark.usefixtures("allow_fall_back_on_numpy")
 @pytest.mark.parametrize(
     "array",

tests/test_sycl_queue.py

@@ -367,6 +367,7 @@ def test_meshgrid(device_x, device_y):
         pytest.param("log1p", [1.0e-10, 1.0, 2.0, 4.0, 7.0]),
         pytest.param("log2", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("max", [1.0, 2.0, 4.0, 7.0]),
+        pytest.param("mean", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("min", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("nancumprod", [1.0, dpnp.nan]),
         pytest.param("nancumsum", [1.0, dpnp.nan]),

tests/test_usm_type.py

@@ -395,6 +395,7 @@ def test_meshgrid(usm_type_x, usm_type_y):
         pytest.param("log2", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("nanprod", [1.0, 2.0, dp.nan]),
         pytest.param("max", [1.0, 2.0, 4.0, 7.0]),
+        pytest.param("mean", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("min", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("negative", [1.0, 0.0, -1.0]),
         pytest.param("positive", [1.0, 0.0, -1.0]),