Implementation of dpnp.mean() (#1431)

vlad-perevezentsev · antonwolfy · antonwolfy · commit 95a24d6d0beb · 2023-06-15T18:47:53.000+02:00
* Reuse dpctl.tensor.sun for dpnp.sum

* Update tests for dpnp.sum

* Fix remarks

* Implementation of dpnp.mean

* Update logic for dpnp.mean function

* add normalize_axis_tuple

* Additional tests for dpnp.mean

* Fix minor remarks

* Add inplace support of divide

* Use inplace divide only for dpnp.inexact types

* Update tests for dpnp.mean

* Skip test_sample.py::TestRandint2::test_bound_float1

* Remove unused import

* Update dtype check

* Update dpnp/dpnp_iface_statistics.py

* Return deleted skips

---------

Co-authored-by: Anton Volkov &lt;antonwolfy@gmail.com&gt;
Co-authored-by: Anton &lt;100830759+antonwolfy@users.noreply.github.com&gt;
diff --git a/dpnp/backend/include/dpnp_iface_fptr.hpp b/dpnp/backend/include/dpnp_iface_fptr.hpp
@@ -218,7 +218,6 @@ enum class DPNPFuncName : size_t
     DPNP_FN_MAXIMUM,                      /**< Used in numpy.maximum() impl  */
     DPNP_FN_MAXIMUM_EXT,                  /**< Used in numpy.maximum() impl , requires extra parameters */
     DPNP_FN_MEAN,                         /**< Used in numpy.mean() impl  */
-    DPNP_FN_MEAN_EXT,                     /**< Used in numpy.mean() impl, requires extra parameters */
     DPNP_FN_MEDIAN,                       /**< Used in numpy.median() impl  */
     DPNP_FN_MEDIAN_EXT,                   /**< Used in numpy.median() impl, requires extra parameters */
     DPNP_FN_MIN,                          /**< Used in numpy.min() impl  */
diff --git a/dpnp/backend/kernels/dpnp_krnl_statistics.cpp b/dpnp/backend/kernels/dpnp_krnl_statistics.cpp
@@ -650,15 +650,6 @@ void (*dpnp_mean_default_c)(void*,
                             const shape_elem_type*,
                             size_t) = dpnp_mean_c<_DataType, _ResultType>;
 
-template <typename _DataType, typename _ResultType>
-DPCTLSyclEventRef (*dpnp_mean_ext_c)(DPCTLSyclQueueRef,
-                                     void*,
-                                     void*,
-                                     const shape_elem_type*,
-                                     const size_t,
-                                     const shape_elem_type*,
-                                     const size_t,
-                                     const DPCTLEventVectorRef) = dpnp_mean_c<_DataType, _ResultType>;
 
 template <typename _DataType, typename _ResultType>
 DPCTLSyclEventRef dpnp_median_c(DPCTLSyclQueueRef q_ref,
@@ -1402,11 +1393,6 @@ void func_map_init_statistics(func_map_t& fmap)
     fmap[DPNPFuncName::DPNP_FN_MEAN][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_mean_default_c<float, float>};
     fmap[DPNPFuncName::DPNP_FN_MEAN][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_mean_default_c<double, double>};
 
-    fmap[DPNPFuncName::DPNP_FN_MEAN_EXT][eft_INT][eft_INT] = {eft_DBL, (void*)dpnp_mean_ext_c<int32_t, double>};
-    fmap[DPNPFuncName::DPNP_FN_MEAN_EXT][eft_LNG][eft_LNG] = {eft_DBL, (void*)dpnp_mean_ext_c<int64_t, double>};
-    fmap[DPNPFuncName::DPNP_FN_MEAN_EXT][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_mean_ext_c<float, float>};
-    fmap[DPNPFuncName::DPNP_FN_MEAN_EXT][eft_DBL][eft_DBL] = {eft_DBL, (void*)dpnp_mean_ext_c<double, double>};
-
     fmap[DPNPFuncName::DPNP_FN_MEDIAN][eft_INT][eft_INT] = {eft_DBL, (void*)dpnp_median_default_c<int32_t, double>};
     fmap[DPNPFuncName::DPNP_FN_MEDIAN][eft_LNG][eft_LNG] = {eft_DBL, (void*)dpnp_median_default_c<int64_t, double>};
     fmap[DPNPFuncName::DPNP_FN_MEDIAN][eft_FLT][eft_FLT] = {eft_FLT, (void*)dpnp_median_default_c<float, float>};
diff --git a/dpnp/dpnp_algo/dpnp_algo.pxd b/dpnp/dpnp_algo/dpnp_algo.pxd
@@ -190,8 +190,6 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_MAX_EXT
         DPNP_FN_MAXIMUM
         DPNP_FN_MAXIMUM_EXT
-        DPNP_FN_MEAN
-        DPNP_FN_MEAN_EXT
         DPNP_FN_MEDIAN
         DPNP_FN_MEDIAN_EXT
         DPNP_FN_MIN
@@ -541,7 +539,6 @@ cpdef dpnp_descriptor dpnp_repeat(dpnp_descriptor array1, repeats, axes=*)
 Statistics functions
 """
 cpdef dpnp_descriptor dpnp_cov(dpnp_descriptor array1)
-cpdef object dpnp_mean(dpnp_descriptor a, axis)
 cpdef dpnp_descriptor dpnp_min(dpnp_descriptor a, axis)
 
 
diff --git a/dpnp/dpnp_algo/dpnp_algo_statistics.pxi b/dpnp/dpnp_algo/dpnp_algo_statistics.pxi
@@ -40,7 +40,6 @@ __all__ += [
     "dpnp_correlate",
     "dpnp_cov",
     "dpnp_max",
-    "dpnp_mean",
     "dpnp_median",
     "dpnp_min",
     "dpnp_nanvar",
@@ -302,152 +301,6 @@ cpdef utils.dpnp_descriptor dpnp_max(utils.dpnp_descriptor x1, axis):
 
     return _dpnp_max(x1, axis_, output_shape)
 
-
-cpdef utils.dpnp_descriptor _dpnp_mean(utils.dpnp_descriptor x1):
-    cdef shape_type_c x1_shape = x1.shape
-    cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(x1.dtype)
-
-    cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_MEAN_EXT, param1_type, param1_type)
-
-    x1_obj = x1.get_array()
-
-    cdef utils.dpnp_descriptor result = utils.create_output_descriptor((1,),
-                                                                       kernel_data.return_type,
-                                                                       None,
-                                                                       device=x1_obj.sycl_device,
-                                                                       usm_type=x1_obj.usm_type,
-                                                                       sycl_queue=x1_obj.sycl_queue)
-
-    result_sycl_queue = result.get_array().sycl_queue
-
-    cdef c_dpctl.SyclQueue q = <c_dpctl.SyclQueue> result_sycl_queue
-    cdef c_dpctl.DPCTLSyclQueueRef q_ref = q.get_queue_ref()
-
-    cdef custom_statistic_1in_1out_func_ptr_t func = <custom_statistic_1in_1out_func_ptr_t > kernel_data.ptr
-
-    # stub for interface support
-    cdef shape_type_c axis
-    cdef Py_ssize_t axis_size = 0
-
-    cdef c_dpctl.DPCTLSyclEventRef event_ref = func(q_ref,
-                                                    x1.get_data(),
-                                                    result.get_data(),
-                                                    x1_shape.data(),
-                                                    x1.ndim,
-                                                    axis.data(),
-                                                    axis_size,
-                                                    NULL)  # dep_events_ref
-
-    with nogil: c_dpctl.DPCTLEvent_WaitAndThrow(event_ref)
-    c_dpctl.DPCTLEvent_Delete(event_ref)
-
-    return result
-
-
-cpdef object dpnp_mean(utils.dpnp_descriptor x1, axis):
-    cdef shape_type_c output_shape
-
-    if axis is None:
-        return _dpnp_mean(x1).get_pyobj()
-
-    cdef long x1_size = x1.size
-    cdef shape_type_c x1_shape = x1.shape
-
-    if x1.dtype == dpnp.float32:
-        res_type = dpnp.float32
-    else:
-        res_type = dpnp.float64
-
-    if x1_size == 0:
-        return dpnp.array([dpnp.nan], dtype=res_type)
-
-    if isinstance(axis, int):
-        axis_ = tuple([axis])
-    else:
-        axis_ = axis
-
-    if axis_ is None:
-        output_shape.push_back(1)
-    else:
-        output_shape = (0, ) * (len(x1_shape) - len(axis_))
-        ind = 0
-        for id, shape_axis in enumerate(x1_shape):
-            if id not in axis_:
-                output_shape[ind] = shape_axis
-                ind += 1
-
-    cdef long prod = 1
-    for i in range(len(output_shape)):
-        if output_shape[i] != 0:
-            prod *= output_shape[i]
-
-    result_array = [None] * prod
-    input_shape_offsets = [None] * len(x1_shape)
-    acc = 1
-
-    for i in range(len(x1_shape)):
-        ind = len(x1_shape) - 1 - i
-        input_shape_offsets[ind] = acc
-        acc *= x1_shape[ind]
-
-    output_shape_offsets = [None] * len(x1_shape)
-    acc = 1
-
-    if axis_ is not None:
-        for i in range(len(output_shape)):
-            ind = len(output_shape) - 1 - i
-            output_shape_offsets[ind] = acc
-            acc *= output_shape[ind]
-            result_offsets = input_shape_offsets[:]  # need copy. not a reference
-        for i in axis_:
-            result_offsets[i] = 0
-
-    for source_idx in range(x1_size):
-
-        # reconstruct x,y,z from linear source_idx
-        xyz = []
-        remainder = source_idx
-        for i in input_shape_offsets:
-            quotient, remainder = divmod(remainder, i)
-            xyz.append(quotient)
-
-        # extract result axis
-        result_axis = []
-        if axis_ is None:
-            result_axis = xyz
-        else:
-            for idx, offset in enumerate(xyz):
-                if idx not in axis_:
-                    result_axis.append(offset)
-
-        # Construct result offset
-        result_offset = 0
-        if axis_ is not None:
-            for i, result_axis_val in enumerate(result_axis):
-                result_offset += (output_shape_offsets[i] * result_axis_val)
-
-        input_elem = input.get_pyobj().item(source_idx)
-        if axis_ is None:
-            if result_array[0] is None:
-                result_array[0] = input_elem
-            else:
-                result_array[0] += input_elem
-        else:
-            if result_array[result_offset] is None:
-                result_array[result_offset] = input_elem
-            else:
-                result_array[result_offset] += input_elem
-
-    del_ = x1_size
-    if axis_ is not None:
-        for i in range(len(x1_shape)):
-            if i not in axis_:
-                del_ = del_ / x1_shape[i]
-    dpnp_array = dpnp.array(result_array, dtype=input.dtype)
-    dpnp_result_array = dpnp.reshape(dpnp_array, output_shape)
-    return dpnp_result_array / del_
-
-
 cpdef utils.dpnp_descriptor dpnp_median(utils.dpnp_descriptor array1):
     cdef shape_type_c x1_shape = array1.shape
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(array1.dtype)
diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py
@@ -794,12 +794,12 @@ def max(self, axis=None, out=None, keepdims=numpy._NoValue, initial=numpy._NoVal
 
         return dpnp.max(self, axis, out, keepdims, initial, where)
 
-    def mean(self, axis=None):
+    def mean(self, axis=None, **kwargs):
         """
         Returns the average of the array elements.
         """
 
-        return dpnp.mean(self, axis)
+        return dpnp.mean(self, axis=axis, **kwargs)
 
     def min(self, axis=None, out=None, keepdims=numpy._NoValue, initial=numpy._NoValue, where=numpy._NoValue):
         """
@@ -1050,7 +1050,7 @@ def transpose(self, *axes):
         array([1, 2, 3, 4])
         >>> a.transpose()
         array([1, 2, 3, 4])
-            
+
         """
 
         ndim = self.ndim
diff --git a/dpnp/dpnp_iface_statistics.py b/dpnp/dpnp_iface_statistics.py
@@ -42,6 +42,7 @@
 
 import numpy
 import dpctl.tensor as dpt
+from numpy.core.numeric import normalize_axis_tuple
 from dpnp.dpnp_algo import *
 from dpnp.dpnp_utils import *
 from dpnp.dpnp_array import dpnp_array
@@ -394,18 +395,23 @@ def max(x1, axis=None, out=None, keepdims=False, initial=None, where=True):
     return call_origin(numpy.max, x1, axis, out, keepdims, initial, where)
 
 
-def mean(x1, axis=None, **kwargs):
+def mean(x, /, *, axis=None, dtype=None, keepdims=False, out=None, where=True):
     """
     Compute the arithmetic mean along the specified axis.
 
     For full documentation refer to :obj:`numpy.mean`.
 
+    Returns
+    -------
+    y : 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.
+
     Limitations
     -----------
-    Input array is supported as :obj:`dpnp.ndarray`.
-    Prameters ``axis`` is supported only with default value ``None``.
-    Keyword arguments ``kwargs`` are currently unsupported.
-    Size of input array is limited by ``x1.size > 0``.
+    Parameters `x` 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.
     Input array data types are limited by supported DPNP :ref:`Data types`.
 
@@ -426,23 +432,52 @@ def mean(x1, axis=None, **kwargs):
     >>> import dpnp as np
     >>> a = np.array([[1, 2], [3, 4]])
     >>> np.mean(a)
-    2.5
-
+    array(2.5)
+    >>> np.mean(a, axis=0)
+    array([2., 3.])
+    >>> np.mean(a, axis=1)
+    array([1.5, 3.5])
     """
 
-    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
-    if x1_desc and not kwargs:
-        if x1_desc.size == 0:
-            pass
-        elif axis is not None:
-            pass
+    if keepdims is not False:
+        pass
+    elif out is not None:
+        pass
+    elif where is not True:
+        pass
+    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:
-            result_obj = dpnp_mean(x1_desc, axis)
-            result = dpnp.convert_single_elem_array_to_scalar(result_obj)
+            new_res_sum = res_sum / del_
+            return new_res_sum.astype(dtype) if new_res_sum.dtype != dtype else new_res_sum
 
-            return result
+        return res_sum.astype(dtype) if res_sum.dtype != dtype else res_sum
 
-    return call_origin(numpy.mean, x1, axis=axis, **kwargs)
+    return call_origin(numpy.mean, x, axis=axis, dtype=dtype, out=out, keepdims=keepdims, where=where)
 
 
 def median(x1, axis=None, out=None, overwrite_input=False, keepdims=False):
diff --git a/tests/skipped_tests.tbl b/tests/skipped_tests.tbl
@@ -989,6 +989,7 @@ tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_1_{siz
 tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_2_{size=4}::test_multinomial
 tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_3_{size=(0,)}::test_multinomial
 tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_4_{size=(1, 0)}::test_multinomial
+tests/third_party/cupy/random_tests/test_sample.py::TestRandint2::test_bound_float1
 tests/third_party/cupy/random_tests/test_sample.py::TestRandint2::test_goodness_of_fit
 tests/third_party/cupy/random_tests/test_sample.py::TestRandint2::test_goodness_of_fit_2
 
@@ -1237,7 +1238,6 @@ tests/third_party/cupy/statistics_tests/test_histogram.py::TestHistogram::test_h
 tests/third_party/cupy/statistics_tests/test_histogram.py::TestHistogram::test_histogram_same_value
 
 tests/third_party/cupy/statistics_tests/test_histogram.py::TestHistogram::test_histogram_weights_mismatch
-tests/third_party/cupy/statistics_tests/test_meanvar.py::TestMeanVar::test_external_mean_axis
 tests/third_party/cupy/statistics_tests/test_meanvar.py::TestNanMeanAdditional::test_nanmean_all_nan
 tests/third_party/cupy/statistics_tests/test_meanvar.py::TestNanMeanAdditional::test_nanmean_float16
 tests/third_party/cupy/statistics_tests/test_meanvar.py::TestNanMeanAdditional::test_nanmean_huge
diff --git a/tests/skipped_tests_gpu.tbl b/tests/skipped_tests_gpu.tbl
@@ -84,6 +84,10 @@ tests/third_party/cupy/math_tests/test_arithmetic.py::TestArithmeticBinary2_para
 tests/third_party/cupy/math_tests/test_arithmetic.py::TestArithmeticBinary2_param_535_{arg1=array([[1, 2, 3],       [4, 5, 6]], dtype=int64), arg2=array([[0, 1, 2],       [3, 4, 5]]), dtype=float64, name='floor_divide', use_dtype=False}::test_binary
 tests/third_party/cupy/math_tests/test_arithmetic.py::TestArithmeticBinary2_param_543_{arg1=array([[1, 2, 3],       [4, 5, 6]], dtype=int64), arg2=array([[0, 1, 2],       [3, 4, 5]], dtype=int64), dtype=float64, name='floor_divide', use_dtype=False}::test_binary
 
+tests/third_party/cupy/math_tests/test_sumprod.py::TestSumprod::test_external_prod_all
+tests/third_party/cupy/math_tests/test_sumprod.py::TestSumprod::test_external_prod_axis
+tests/third_party/cupy/math_tests/test_sumprod.py::TestSumprod::test_prod_all
+tests/third_party/cupy/math_tests/test_sumprod.py::TestSumprod::test_prod_axis
 tests/third_party/cupy/math_tests/test_sumprod.py::TestSumprod::test_sum_out
 tests/third_party/cupy/math_tests/test_sumprod.py::TestSumprod::test_sum_out_wrong_shape
 tests/third_party/cupy/math_tests/test_sumprod.py::TestCumprod::test_cumprod_1dim
@@ -250,7 +254,6 @@ tests/third_party/cupy/random_tests/test_distributions.py::TestDistributionsMult
 tests/third_party/cupy/random_tests/test_distributions.py::TestDistributionsMultivariateNormal_param_3_{d=4, shape=(3, 2)}::test_normal
 
 tests/third_party/cupy/statistics_tests/test_correlation.py::TestCov::test_cov_empty
-tests/third_party/cupy/statistics_tests/test_meanvar.py::TestMeanVar::test_external_mean_axis
 
 tests/third_party/intel/test_zero_copy_test1.py::test_dpnp_interaction_with_dpctl_memory
 tests/test_arraymanipulation.py::TestHstack::test_generator
@@ -1075,6 +1078,7 @@ tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_1_{siz
 tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_2_{size=4}::test_multinomial
 tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_3_{size=(0,)}::test_multinomial
 tests/third_party/cupy/random_tests/test_sample.py::TestMultinomial_param_4_{size=(1, 0)}::test_multinomial
+tests/third_party/cupy/random_tests/test_sample.py::TestRandint2::test_bound_float1
 tests/third_party/cupy/random_tests/test_sample.py::TestRandint2::test_goodness_of_fit
 tests/third_party/cupy/random_tests/test_sample.py::TestRandint2::test_goodness_of_fit_2
 tests/third_party/cupy/random_tests/test_sample.py::TestRandomIntegers2::test_bound_1
diff --git a/tests/test_mathematical.py b/tests/test_mathematical.py
diff --git a/tests/third_party/cupy/statistics_tests/test_meanvar.py b/tests/third_party/cupy/statistics_tests/test_meanvar.py
