Repository: numba/numba

Issue:
Overflow in aggregation functions when working with smaller types
The implementations of Numpy aggregation functions have accumulators that overflow too quickly when operating on arrays with small types:

``` python
import numpy as np
import numba

@numba.jit
def numba_mean(x):
    return x.mean()

a = np.full(shape=100, fill_value=16000, dtype=np.int16)

print('Numpy:', np.mean(a), type(np.mean(a)))
print('Numba:', numba_mean(a), type(numba_mean(a)))
```

produces output:

```
Numpy: 16000.0 <class 'numpy.float64'>
Numba: 271.36 <class 'float'>
```

We need to ensure that the accumulators are using int64 or float64 for all overflowable aggregation functions (so not `min()` and `max()`).



Discussion:

Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -436,7 +436,7 @@ def array_sum_impl(arr):
         return c
 
     return context.compile_internal(builder, array_sum_impl, sig, args, 
-                                    locals=dict(c=arrty.dtype))
+                                    locals=dict(c=sig.return_type))
 
 
 @builtin
@@ -452,7 +452,7 @@ def array_prod_impl(arr):
         return c
 
     return context.compile_internal(builder, array_prod_impl, sig, args,
-                                    locals=dict(c=arrty.dtype))
+                                    locals=dict(c=sig.return_type))
 
 
 @builtin
@@ -461,10 +461,16 @@ def array_prod_impl(arr):
 def array_mean(context, builder, sig, args):
     [arrty] = sig.args
 
-    def array_mean_impl(arry):
-        return arry.sum() / arry.size
+    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 = 0
+        for v in arr.flat:
+            c += v
+        return c / arr.size
 
-    return context.compile_internal(builder, array_mean_impl, sig, args)
+    return context.compile_internal(builder, array_mean_impl, sig, args,
+                                    locals=dict(c=sig.return_type))
 
 
 @builtin
diff --git a/numba/typing/builtins.py b/numba/typing/builtins.py
--- a/numba/typing/builtins.py
+++ b/numba/typing/builtins.py
@@ -485,8 +485,15 @@ def generic_homog(self, args, kws):
     assert not kws
     return signature(self.this.dtype, recvr=self.this)
 
+def generic_expand(self, args, kws):
+    if isinstance(self.this.dtype, types.Integer):
+        if self.this.dtype.signed:
+            return signature(types.int64, recvr=self.this)
+        else:
+            return signature(types.uint64, recvr=self.this)
+    return signature(self.this.dtype, recvr=self.this)
 
-def generic_hetero(self, args, kws):
+def generic_hetero_real(self, args, kws):
     assert not args
     assert not kws
     if self.this.dtype in types.integer_domain:
@@ -508,14 +515,18 @@ def array_attribute_attachment(self, ary):
     setattr(ArrayAttribute, "resolve_" + name, array_attribute_attachment)
 
 # Functions that return the same type as the array
-for fName in ["sum", "prod", "min", "max"]:
-    install_array_method(fName, generic_homog)
+for fname in ["min", "max"]:
+    install_array_method(fname, generic_homog)
+
+# Functions that return a type as large as possible, to avoid overflows
+for fname in ["sum", "prod"]:
+    install_array_method(fname, generic_expand)
 
 # Functions that require integer arrays get promoted to float64 return
 for fName in ["mean", "var", "std"]:
-    install_array_method(fName, generic_hetero)
+    install_array_method(fName, generic_hetero_real)
 
-# Functions that return an index (int64)
+# Functions that return an index (intp)
 install_array_method("argmin", generic_index)
 install_array_method("argmax", generic_index)
 
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -184,16 +184,27 @@ class typing_class(Numpy_rules_ufunc):
 # -----------------------------------------------------------------------------
 # Install global reduction functions
 
-# Functions where domain and range are the same format
+# Functions where input domain and output domain are the same
 class Numpy_homogenous_reduction(AbstractTemplate):
     def generic(self, args, kws):
         assert not kws
         [arr] = args
         return signature(arr.dtype, arr)
 
-
 # Functions where domain and range are possibly different formats
-class Numpy_heterogenous_reduction(AbstractTemplate):
+class Numpy_expanded_reduction(AbstractTemplate):
+    def generic(self, args, kws):
+        assert not kws
+        [arr] = args
+        if isinstance(arr.dtype, types.Integer):
+            if arr.dtype.signed:
+                return signature(types.int64, arr)
+            else:
+                return signature(types.uint64, arr)
+        else:
+            return signature(arr.dtype, arr)
+
+class Numpy_heterogenous_reduction_real(AbstractTemplate):
     def generic(self, args, kws):
         assert not kws
         [arr] = args
@@ -202,7 +213,6 @@ def generic(self, args, kws):
         else:
             return signature(arr.dtype, arr)
 
-
 class Numpy_index_reduction(AbstractTemplate):
     def generic(self, args, kws):
         assert not kws
@@ -213,16 +223,19 @@ def generic(self, args, kws):
 def _numpy_reduction(fname, rClass):
     npyfn = getattr(numpy, fname)
     cls = type("Numpy_reduce_{0}".format(npyfn), (rClass,), dict(key=npyfn))
-    semiBound = lambda self, mod: types.Function(cls) 
+    semiBound = lambda self, mod: types.Function(cls)
     setattr(NumpyModuleAttribute, "resolve_{0}".format(fname), semiBound)
 
-for func in ['sum', 'prod', 'min', 'max']:
+for func in ['min', 'max']:
     _numpy_reduction(func, Numpy_homogenous_reduction)
 
+for func in ['sum', 'prod']:
+    _numpy_reduction(func, Numpy_expanded_reduction)
+
 for func in ['mean', 'var', 'std']:
-    _numpy_reduction(func, Numpy_heterogenous_reduction)
+    _numpy_reduction(func, Numpy_heterogenous_reduction_real)
 
-for func in ["argmin", "argmax"]:
+for func in ['argmin', 'argmax']:
     _numpy_reduction(func, Numpy_index_reduction)
 
 


Test patch:
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -284,6 +284,41 @@ def test_array_prod_global(self):
 
         np.testing.assert_allclose(np.prod(arr), cfunc(arr))
 
+    def check_aggregation_magnitude(self, pyfunc):
+        """
+        Check that integer overflows are avoided (issue #931).
+        """
+        # Overflows are always avoided here (ints are cast either to int64
+        # or float64).
+        arr = np.arange(10, dtype='int16') + 60000
+        npr, nbr = run_comparative(pyfunc, arr)
+        self.assertPreciseEqual(npr, nbr)
+        # Overflows are avoided for functions returning floats here.
+        arr = np.arange(10, dtype='int64') + 2**63
+        npr, nbr = run_comparative(pyfunc, arr)
+        self.assertPreciseEqual(npr, nbr)
+
+    def test_sum_magnitude(self):
+        self.check_aggregation_magnitude(array_sum)
+        self.check_aggregation_magnitude(array_sum_global)
+
+    def test_prod_magnitude(self):
+        self.check_aggregation_magnitude(array_prod)
+        self.check_aggregation_magnitude(array_prod_global)
+
+    def test_mean_magnitude(self):
+        self.check_aggregation_magnitude(array_mean)
+        self.check_aggregation_magnitude(array_mean_global)
+
+    def test_var_magnitude(self):
+        self.check_aggregation_magnitude(array_var)
+        self.check_aggregation_magnitude(array_var_global)
+
+    def test_std_magnitude(self):
+        self.check_aggregation_magnitude(array_std)
+        self.check_aggregation_magnitude(array_std_global)
+
+
 # These form a testing product where each of the combinations are tested
 reduction_funcs = [array_sum, array_sum_global,
                    array_prod, array_prod_global,