[WIP] Reimplementing np.mean with overload and adding axis parameter

docs/source/reference/numpysupported.rst

@@ -182,7 +182,6 @@ The following methods of Numpy arrays are supported in their basic form
 * :meth:`~numpy.ndarray.cumprod`
 * :meth:`~numpy.ndarray.cumsum`
 * :meth:`~numpy.ndarray.max`
-* :meth:`~numpy.ndarray.mean`
 * :meth:`~numpy.ndarray.min`
 * :meth:`~numpy.ndarray.nonzero`
 * :meth:`~numpy.ndarray.prod`
@@ -287,6 +286,14 @@ Reductions
 The following reduction functions are supported:
 
 * :func:`numpy.diff` (only the 2 first arguments)
+* :func:`numpy.mean` (with or without the ``axis`` argument)
+
+  * ``axis`` only supports ``integer`` values.
+  * If the ``axis`` argument is a compile-time constant, all valid values
+    are supported.
+    An out-of-range value will result in a ``LoweringError`` at compile-time.
+  * If the ``axis`` argument is not a compile-time constant, only values
+    from 0 to 3 are supported.
 * :func:`numpy.median` (only the first argument)
 * :func:`numpy.nancumprod` (only the first argument, requires NumPy >= 1.12))
 * :func:`numpy.nancumsum` (only the first argument, requires NumPy >= 1.12))

numba/targets/arraymath.py

@@ -399,22 +399,52 @@ def array_cumprod_impl(arr):
     return impl_ret_new_ref(context, builder, sig.return_type, res)
 
 
-@lower_builtin(np.mean, types.Array)
-@lower_builtin("array.mean", types.Array)
-def array_mean(context, builder, sig, args):
-    zero = sig.return_type(0)
+@register_jitable
+def sum_array(arr, dtype):
+    return np.sum(arr, dtype=dtype)
 
-    def array_mean_impl(arr):
-        # Can't use the naive `arr.sum() / arr.size`, as it would return
-        # a wrong result on integer sum overflow.
-        c = zero
-        for v in np.nditer(arr):
-            c += v.item()
-        return c / arr.size
 
-    res = context.compile_internal(builder, array_mean_impl, sig, args,
-                                   locals=dict(c=sig.return_type))
-    return impl_ret_untracked(context, builder, sig.return_type, res)
+@register_jitable
+def sum_array_axis(arr, axis, dtype):
+    return np.sum(arr, axis=axis, dtype=dtype)
+
+
+@overload(np.mean)
+@overload_method(types.Array, 'mean')
+def array_mean(arr, axis=None):
+    if isinstance(arr, types.Array):
+        # determine accumulator type
+        if isinstance(arr.dtype, (types.Integer, types.Boolean)):
+            ret_dtype = np.float64
+        elif isinstance(arr.dtype, (types.Float, types.Complex)):
+            ret_dtype = arr.dtype
+        else:
+            raise TypeError(("np.mean is not supported on {} arrays. "
+                            "It supports boolean, integer, float "
+                            "and complex arrays").format(arr.dtype))
+        # dispatch based on whether there's an axis parameter and its type
+        if axis is None:
+
+            def mean_impl(arr, axis=None):
+                return sum_array(arr, dtype=ret_dtype) / arr.size
+
+            return mean_impl
+        elif isinstance(axis, types.Integer):
+
+            def mean_impl(arr, axis=None):
+                if axis >= arr.ndim:
+                    raise ValueError("'axis' entry is out of bounds")
+                return sum_array_axis(arr, axis=axis, dtype=ret_dtype) / arr.shape[axis]
+
+            return mean_impl
+        elif isinstance(axis, types.IntegerLiteral):
+            if axis.literal_value >= arr.ndim:
+                raise ValueError("'axis' entry is out of bounds")
+
+            def mean_impl(arr, axis=None):
+                return sum_array_axis(arr, axis=axis, dtype=ret_dtype) / arr.shape[axis]
+
+            return mean_impl
 
 
 @lower_builtin(np.var, types.Array)

numba/tests/test_array_methods.py

@@ -179,6 +179,12 @@ def array_sum_const_axis_neg_one(a, axis):
     # "const_multi" variant would raise errors
     return a.sum(axis=-1)
 
+def array_mean(a):
+    return a.mean()
+
+def array_mean_axis(a, axis):
+    return a.mean(axis)
+
 def array_cumsum(a, *args):
     return a.cumsum(*args)
 
@@ -784,7 +790,8 @@ def test_sum(self):
         pyfunc = array_sum
         cfunc = jit(nopython=True)(pyfunc)
         all_dtypes = [np.float64, np.float32, np.int64, np.int32,
-                      np.complex64, np.complex128, np.uint32, np.uint64, np.timedelta64]
+                      np.complex64, np.complex128, np.uint32, np.uint64,
+                      np.timedelta64, np.bool_]
         all_test_arrays = [
             [np.ones((7, 6, 5, 4, 3), arr_dtype),
              np.ones(1, arr_dtype),
@@ -1006,6 +1013,39 @@ def foo(arr):
         # Just check for the "out of bounds" phrase in it.
         self.assertIn("out of bounds", str(raises.exception))
 
+    def test_mean_axis(self):
+        """   tests np.mean with and without axis parameter
+        """
+        pyfunc = array_mean
+        cfunc = jit(nopython=True)(pyfunc)
+        pyfunc_axis = array_mean_axis
+        cfunc_axis = jit(nopython=True)(pyfunc_axis)
+        # a complete list
+        # all_dtypes = [np.float64, np.float32, np.int64, np.int32, np.uint32,
+        #               np.uint64, np.complex64, np.complex128]
+        # a reduced list to save test execution time
+        all_dtypes = [np.float32, np.int32, np.uint32, np.complex64, np.bool_]
+        all_dtypes = [np.float64]
+        all_test_arrays = [np.ones((7, 6, 5, 4, 3), arr_dtype) for arr_dtype
+                           in all_dtypes]
+
+        for arr in all_test_arrays:
+            with self.subTest("no axis - dtype: {}".format(arr.dtype)):
+                self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+            # with self.subTest("axis 0 as integer literal - dtype: {}".format(arr.dtype)):
+            #     self.assertPreciseEqual(pyfunc_axis(arr, 0), cfunc_axis(arr, 0))
+            # with self.subTest("axis 1 as integer variable - dtype: {}".format(arr.dtype)):
+            #     axis = 1
+            #     self.assertPreciseEqual(pyfunc_axis(arr, axis), cfunc_axis(arr, axis))
+            # with self.subTest("axis 2  as integer variable - dtype: {}".format(arr.dtype)):
+            #     axis = 2
+            #     self.assertPreciseEqual(pyfunc_axis(arr, axis), cfunc_axis(arr, axis))
+            # with self.subTest("axis -1 as integer literal- dtype: {}".format(arr.dtype)):
+            #     # axis -1 is only supported for IntegerLiterals
+            #     pyfunc2 = lambda x: x.mean(axis=-1)
+            #     cfunc2 = jit(nopython=True)(pyfunc2)
+            #     self.assertPreciseEqual(pyfunc2(arr), cfunc2(arr))
+
     def test_cumsum(self):
         pyfunc = array_cumsum
         cfunc = jit(nopython=True)(pyfunc)

numba/typing/arraydecl.py

@@ -730,8 +730,8 @@ def array_attribute_attachment(self, ary):
     install_array_method(fname, generic_expand_cumulative)
 
 # Functions that require integer arrays get promoted to float64 return
-for fName in ["mean"]:
-    install_array_method(fName, generic_hetero_real)
+# for fName in ["mean"]:
+#     install_array_method(fName, generic_hetero_real)
 
 # var and std by definition return in real space and int arrays
 # get promoted to float64 return

numba/typing/npydecl.py

@@ -414,7 +414,7 @@ def _numpy_redirect(fname):
                dict(key=numpy_function, method_name=fname))
     infer_global(numpy_function, types.Function(cls))
 
-for func in ['min', 'max', 'sum', 'prod', 'mean', 'var', 'std',
+for func in ['min', 'max', 'sum', 'prod', 'var', 'std',  #'mean',
              'cumsum', 'cumprod', 'argmin', 'argmax', 'argsort',
              'nonzero', 'ravel']:
     _numpy_redirect(func)