Repository: numba/numba
Pull request: Allow dtype input argument in np.sum
closes #4220



Problem statement:
support for 'dtype' keyword argument in np.sum()
## Feature request: Add `dtype` keyword argument to `np.sum()`

I am working on a scientific toolbox that relies on numpy and scipy and requires decent numeric precision. To parallelize core parts, I considered using Numba as an _optional dependency_ with a fallback to native numpy. This seems doable and elegant except for the following:

The existing code relies heavily on `np.mean()` with the `axis` and `dtype` keyword arguments. `dtype=np.float64` is needed to get enough precision in certain situations.

* This does _not_ cast (copy to new dtype) the input array, only the temporary variable is of a higher precision.
* I managed to implement this in numba with a manual for loop, casting _within_ the loop
* The problem arises when using arrays of higher dimension. A specific function has to be written for every ndim-to-axis combination. As far as i understand, one cannot tell numba to check if-conditions for dimensionality during compilation, to avoid `tuple index out of range`.

The script below illustrates the precision issue:

```
import numpy as np
import numpy.random
import time

try:
    from numba import jit, njit, prange, gdb
    # raise ImportError # uncomment to toggle numba
except ImportError:
    print('Numba not available, skipping parallelization')
    # replace numba functions if numba not available:
    # we only use jit and prange
    # helper needed for decorators with kwargs
    def parametrized(dec):
        def layer(*args, **kwargs):
            def repl(f):
                return dec(f, *args, **kwargs)
            return repl
        return layer

    @parametrized
    def jit(func, **kwargs):
        return func
    def prange(*args):
        return range(*args)

ftype = np.float64


@jit(nopython=True, parallel=False, fastmath=True)
def mean_lowprec(a):
    return np.sum(a)/(2*a.shape[1])

@jit(nopython=True, parallel=False, fastmath=True)
def mean_copy(a):
    # copies in ram, factor two memory required
    b = a.astype(ftype)
    return np.sum(b)/(2*a.shape[1])

@jit(nopython=True, parallel=False, fastmath=True)
def mean_2d_reduce(a):
    total = ftype(0)
    for i in range(a.shape[0]):
        for j in range(a.shape[1]):
            total += ftype(a[i,j])

    return total/ftype(a.shape[0]*a.shape[1])


a = np.zeros((2, 512*512), dtype=np.float32)
a[0, :] = 1.0
a[1, :] = 0.1

# compile first call
mean_lowprec(a)
mean_copy(a)
mean_2d_reduce(a)

time_beg = time.time()
print(f'np.sum f4:           {np.sum(a)/(2*a.shape[1])}')
print(f'{(time.time()-time_beg)*1000 :.2f} ms\n')

time_beg = time.time()
print(f'np.sum f8:           {np.sum(a, dtype=ftype)/(2*a.shape[1])}')
print(f'{(time.time()-time_beg)*1000 :.2f} ms\n')

time_beg = time.time()
print(f'mean_lowprec(a):     {mean_lowprec(a)}')
print(f'{(time.time()-time_beg)*1000 :.2f} ms\n')

time_beg = time.time()
print(f'mean_copy(a):        {mean_copy(a)}')
print(f'{(time.time()-time_beg)*1000 :.2f} ms\n')

time_beg = time.time()
print(f'mean_2d_reduce(a):   {mean_2d_reduce(a)}')
print(f'{(time.time()-time_beg)*1000 :.2f} ms\n')
```

Results in
```
np.sum f4:           0.5499992370605469
0.49 ms

np.sum f8:           0.5500000007450581
0.55 ms

mean_lowprec(a):     0.5498046875
0.10 ms

mean_copy(a):        0.5500000007450581
0.84 ms

mean_2d_reduce(a):   0.5500000007450581
0.15 ms
```

While the last option seems a good candidate performance/memory wise, I failed to come up with a solution that offers a one-wrapper-function fits-all to mimic the native `numpy.mean()`.

Hence, I believe support for the `dtype` argument in numba would be a great feature!
I would greatly appreciate ideas for a workaround or some background on the details I am missing. 

Thanks,
Paul

Hints from issue discussion:
Summarizing:

- [ ] feature request for `np.mean()` with `dtype` argument.
- [ ] additional feature request for `np.sum()` with `dtype` argument (because `np.mean` calls `np.sum`).

The ability to do `if my_array.ndim == 2` relates to the dead branch pruning feature, but it doesn't do work in this case.  @stuartarchibald, i think this case will require type-inference info in branch pruning.
@sklam think the semantic constant rewrites https://github.com/numba/numba/pull/4113 will replace the `my_array.ndim` in `if my_array.ndim == 2` with the constant  value, e.g. were it 2, the condition would end up `if 2 == 2` which would prune fine?

Commit message: fix for scalar results

Patch to review:
diff --git a/docs/source/reference/numpysupported.rst b/docs/source/reference/numpysupported.rst
index f13a810c6e2..a9760dea164 100644
--- a/docs/source/reference/numpysupported.rst
+++ b/docs/source/reference/numpysupported.rst
@@ -210,13 +210,21 @@ The following methods of Numpy arrays are supported:
 * :meth:`~numpy.ndarray.repeat` (no axis argument)
 * :meth:`~numpy.ndarray.reshape` (only the 1-argument form)
 * :meth:`~numpy.ndarray.sort` (without arguments)
-* :meth:`~numpy.ndarray.sum` (with or without the ``axis`` argument. 
-  ``axis`` only supports ``integer`` values)
+* :meth:`~numpy.ndarray.sum` (with or without the ``axis`` and/or ``dtype`` arguments.)
 
+  * ``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.
     An out-of-range value will result in a runtime exception.
+  * All numeric ``dtypes`` are supported as ``dtype`` parameter. ``timedelta`` arrays can
+    be used as input arrays but ``timedelta`` is not supported as ``dtype`` parameter.
+  * When ``dtype`` is given, it determines the type of the internal accumulator. When it is not,
+    the selection is made automatically based on the input array's ``dtype``, mostly following the same
+    rules as NumPy. However, on 64-bit Windows, Numba uses a 64-bit accumulator for integer inputs
+    (``int64`` for ``int32`` inputs and ``uint64`` for ``uint32`` inputs), while NumPy would use a 32-bit
+    accumulator in those cases.
+
 
 * :meth:`~numpy.ndarray.transpose`
 * :meth:`~numpy.ndarray.view` (only the 1-argument form)
diff --git a/numba/targets/arraymath.py b/numba/targets/arraymath.py
index 5135254dda1..d6086f0f4bc 100644
--- a/numba/targets/arraymath.py
+++ b/numba/targets/arraymath.py
@@ -177,50 +177,8 @@ def _array_sum_axis_nop(arr, v):
     return arr
 
 
-@lower_builtin(np.sum, types.Array, types.intp)
-@lower_builtin(np.sum, types.Array, types.IntegerLiteral)
-@lower_builtin("array.sum", types.Array, types.intp)
-@lower_builtin("array.sum", types.Array, types.IntegerLiteral)
-def array_sum_axis(context, builder, sig, args):
-    """
-    The third parameter to gen_index_tuple that generates the indexing
-    tuples has to be a const so we can't just pass "axis" through since
-    that isn't const.  We can check for specific values and have
-    different instances that do take consts.  Supporting axis summation
-    only up to the fourth dimension for now.
-    """
-    # typing/arraydecl.py:sum_expand defines the return type for sum with axis.
-    # It is one dimension less than the input array.
-
-    retty = sig.return_type
-    zero = getattr(retty, 'dtype', retty)(0)
-    # if the return is scalar in type then "take" the 0th element of the
-    # 0d array accumulator as the return value
-    if getattr(retty, 'ndim', None) is None:
-        op = np.take
-    else:
-        op = _array_sum_axis_nop
-    [ty_array, ty_axis] = sig.args
-    is_axis_const = False
-    const_axis_val = 0
-    if isinstance(ty_axis, types.Literal):
-        # this special-cases for constant axis
-        const_axis_val = ty_axis.literal_value
-        # fix negative axis
-        if const_axis_val < 0:
-            const_axis_val = ty_array.ndim + const_axis_val
-        if const_axis_val < 0 or const_axis_val > ty_array.ndim:
-            raise ValueError("'axis' entry is out of bounds")
-
-        ty_axis = context.typing_context.resolve_value_type(const_axis_val)
-        axis_val = context.get_constant(ty_axis, const_axis_val)
-        # rewrite arguments
-        args = args[0], axis_val
-        # rewrite sig
-        sig = sig.replace(args=[ty_array, ty_axis])
-        is_axis_const = True
-
-    def array_sum_impl_axis(arr, axis):
+def gen_sum_axis_impl(is_axis_const, const_axis_val, op, zero):
+    def inner(arr, axis):
         ndim = arr.ndim
 
         if not is_axis_const:
@@ -268,8 +226,129 @@ def array_sum_impl_axis(arr, axis):
                 elif axis == 3:
                     index_tuple4 = _gen_index_tuple(arr.shape, axis_index, 3)
                     result += arr[index_tuple4]
-
         return op(result, 0)
+    return inner
+
+
+@lower_builtin(np.sum, types.Array, types.intp, types.DTypeSpec)
+@lower_builtin(np.sum, types.Array, types.IntegerLiteral, types.DTypeSpec)
+@lower_builtin("array.sum", types.Array, types.intp, types.DTypeSpec)
+@lower_builtin("array.sum", types.Array, types.IntegerLiteral, types.DTypeSpec)
+def array_sum_axis_dtype(context, builder, sig, args):
+    """
+    The third parameter to gen_index_tuple that generates the indexing
+    tuples has to be a const so we can't just pass "axis" through since
+    that isn't const.  We can check for specific values and have
+    different instances that do take consts.  Supporting axis summation
+    only up to the fourth dimension for now.
+    """
+    # typing/arraydecl.py:sum_expand defines the return type for sum with axis.
+    # It is one dimension less than the input array.
+
+    retty = sig.return_type
+    zero = getattr(retty, 'dtype', retty)(0)
+    # if the return is scalar in type then "take" the 0th element of the
+    # 0d array accumulator as the return value
+    if getattr(retty, 'ndim', None) is None:
+        op = np.take
+    else:
+        op = _array_sum_axis_nop
+    [ty_array, ty_axis, ty_dtype] = sig.args
+    is_axis_const = False
+    const_axis_val = 0
+    if isinstance(ty_axis, types.Literal):
+        # this special-cases for constant axis
+        const_axis_val = ty_axis.literal_value
+        # fix negative axis
+        if const_axis_val < 0:
+            const_axis_val = ty_array.ndim + const_axis_val
+        if const_axis_val < 0 or const_axis_val > ty_array.ndim:
+            raise ValueError("'axis' entry is out of bounds")
+
+        ty_axis = context.typing_context.resolve_value_type(const_axis_val)
+        axis_val = context.get_constant(ty_axis, const_axis_val)
+        # rewrite arguments
+        args = args[0], axis_val, args[2]
+        # rewrite sig
+        sig = sig.replace(args=[ty_array, ty_axis, ty_dtype])
+        is_axis_const = True
+
+    gen_impl = gen_sum_axis_impl(is_axis_const, const_axis_val, op, zero)
+    compiled = register_jitable(gen_impl)
+
+    def array_sum_impl_axis(arr, axis, dtype):
+        return compiled(arr, axis)
+
+    res = context.compile_internal(builder, array_sum_impl_axis, sig, args)
+    return impl_ret_new_ref(context, builder, sig.return_type, res)
+
+
+@lower_builtin(np.sum, types.Array,  types.DTypeSpec)
+@lower_builtin(np.sum, types.Array,  types.DTypeSpec)
+@lower_builtin("array.sum", types.Array, types.DTypeSpec)
+@lower_builtin("array.sum", types.Array, types.DTypeSpec)
+def array_sum_dtype(context, builder, sig, args):
+    zero = sig.return_type(0)
+
+    def array_sum_impl(arr, dtype):
+        c = zero
+        for v in np.nditer(arr):
+            c += v.item()
+        return c
+
+    res = context.compile_internal(builder, array_sum_impl, sig, args,
+                                   locals=dict(c=sig.return_type))
+    return impl_ret_borrowed(context, builder, sig.return_type, res)
+
+
+@lower_builtin(np.sum, types.Array, types.intp)
+@lower_builtin(np.sum, types.Array, types.IntegerLiteral)
+@lower_builtin("array.sum", types.Array, types.intp)
+@lower_builtin("array.sum", types.Array, types.IntegerLiteral)
+def array_sum_axis(context, builder, sig, args):
+    """
+    The third parameter to gen_index_tuple that generates the indexing
+    tuples has to be a const so we can't just pass "axis" through since
+    that isn't const.  We can check for specific values and have
+    different instances that do take consts.  Supporting axis summation
+    only up to the fourth dimension for now.
+    """
+    # typing/arraydecl.py:sum_expand defines the return type for sum with axis.
+    # It is one dimension less than the input array.
+
+    retty = sig.return_type
+    zero = getattr(retty, 'dtype', retty)(0)
+    # if the return is scalar in type then "take" the 0th element of the
+    # 0d array accumulator as the return value
+    if getattr(retty, 'ndim', None) is None:
+        op = np.take
+    else:
+        op = _array_sum_axis_nop
+    [ty_array, ty_axis] = sig.args
+    is_axis_const = False
+    const_axis_val = 0
+    if isinstance(ty_axis, types.Literal):
+        # this special-cases for constant axis
+        const_axis_val = ty_axis.literal_value
+        # fix negative axis
+        if const_axis_val < 0:
+            const_axis_val = ty_array.ndim + const_axis_val
+        if const_axis_val < 0 or const_axis_val > ty_array.ndim:
+            raise ValueError("'axis' entry is out of bounds")
+
+        ty_axis = context.typing_context.resolve_value_type(const_axis_val)
+        axis_val = context.get_constant(ty_axis, const_axis_val)
+        # rewrite arguments
+        args = args[0], axis_val
+        # rewrite sig
+        sig = sig.replace(args=[ty_array, ty_axis])
+        is_axis_const = True
+
+    gen_impl = gen_sum_axis_impl(is_axis_const, const_axis_val, op, zero)
+    compiled = register_jitable(gen_impl)
+
+    def array_sum_impl_axis(arr, axis):
+        return compiled(arr, axis)
 
     res = context.compile_internal(builder, array_sum_impl_axis, sig, args)
     return impl_ret_new_ref(context, builder, sig.return_type, res)
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
index d2704817c20..a199b8cd4ec 100644
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -148,9 +148,18 @@ def array_itemset(a, v):
 def array_sum(a, *args):
     return a.sum(*args)
 
-def array_sum_kws(a, axis):
+def array_sum_axis_kws(a, axis):
     return a.sum(axis=axis)
 
+def array_sum_dtype_kws(a, dtype):
+    return a.sum(dtype=dtype)
+
+def array_sum_axis_dtype_kws(a, dtype, axis):
+    return a.sum(axis=axis, dtype=dtype)
+
+def array_sum_axis_dtype_pos(a, a1, a2):
+    return a.sum(a1, a2)
+
 def array_sum_const_multi(arr, axis):
     # use np.sum with different constant args multiple times to check
     # for internal compile cache to see if constant-specialization is
@@ -492,7 +501,6 @@ def check_arr(arr):
         arr = np.array([0]).reshape(())
         check_arr(arr)
 
-
     def test_array_transpose(self):
         self.check_layout_dependent_func(array_transpose)
 
@@ -531,7 +539,7 @@ def check(shape):
     def test_np_frombuffer_allocated(self):
         self.check_np_frombuffer_allocated(np_frombuffer_allocated)
 
-    def test_np_frombuffer_allocated(self):
+    def test_np_frombuffer_allocated2(self):
         self.check_np_frombuffer_allocated(np_frombuffer_allocated_dtype)
 
     def check_nonzero(self, pyfunc):
@@ -779,8 +787,24 @@ def test_sum(self):
         # OK
         self.assertPreciseEqual(pyfunc(a, 0), cfunc(a, 0))
 
-    def test_sum_kws(self):
-        pyfunc = array_sum_kws
+    def test_sum2(self):
+        pyfunc = array_sum
+        cfunc = jit(nopython=True)(pyfunc)
+        all_dtypes = [np.float64, np.float32, np.int64, np.int32,
+                      np.complex64, np.uint32, np.uint64, np.timedelta64]
+        all_test_arrays = [
+            [np.ones((7, 6, 5, 4, 3), arr_dtype),
+             np.ones(1, arr_dtype),
+             np.ones((7, 3), arr_dtype) * -5]
+            for arr_dtype in all_dtypes]
+
+        for arr_list in all_test_arrays:
+            for arr in arr_list:
+                with self.subTest("Test np.sum with {} input ".format(arr.dtype)):
+                    self.assertPreciseEqual(pyfunc(arr), cfunc(arr))
+
+    def test_sum_axis_kws(self):
+        pyfunc = array_sum_axis_kws
         cfunc = jit(nopython=True)(pyfunc)
         # OK
         a = np.ones((7, 6, 5, 4, 3))
@@ -788,9 +812,170 @@ def test_sum_kws(self):
         # OK
         self.assertPreciseEqual(pyfunc(a, axis=2), cfunc(a, axis=2))
 
+    def test_sum_axis_kws2(self):
+        pyfunc = array_sum_axis_kws
+        cfunc = jit(nopython=True)(pyfunc)
+        all_dtypes = [np.float64, np.float32, np.int64, np.uint64, np.complex64]
+        # timedelta test cannot be enabled until issue #4540 is fixed
+        # all_dtypes += [np.timedelta64]
+        all_test_arrays = [
+            [np.ones((7, 6, 5, 4, 3), arr_dtype),
+             np.ones(1, arr_dtype),
+             np.ones((7, 3), arr_dtype) * -5]
+            for arr_dtype in all_dtypes]
+
+        for arr_list in all_test_arrays:
+            for arr in arr_list:
+                for axis in (0, 1, 2):
+                    if axis > len(arr.shape)-1:
+                        continue
+                    with self.subTest("Testing np.sum(axis) with {} "
+                                      "input ".format(arr.dtype)):
+                        self.assertPreciseEqual(pyfunc(arr, axis=axis),
+                                                cfunc(arr, axis=axis))
+
+    def test_sum_axis_kws3(self):
+        """ uint32 and int32 must be tested separately because Numpy's current
+        behaviour is different in 64bits Windows (accumulates as int32)
+        and 64bits Linux (accumulates as int64), while Numba has decided to always
+        accumulate as int64, when the OS is 64bits. No testing has been done
+        for behaviours in 32 bits platforms.
+
+        :return:
+        """
+        pyfunc = array_sum_axis_kws
+        cfunc = jit(nopython=True)(pyfunc)
+        all_dtypes = [np.int32, np.uint32]
+        # expected return dtypes in Numba
+        out_dtypes = {np.dtype('int32'): np.int64, np.dtype('uint32'): np.uint64,
+                      np.dtype('int64'): np.int64}
+        # timedelta test cannot be enabled until issue #4540 is fixed
+        # all_dtypes += [np.timedelta64]
+        all_test_arrays = [
+            [np.ones((7, 6, 5, 4, 3), arr_dtype),
+             np.ones(1, arr_dtype),
+             np.ones((7, 3), arr_dtype) * -5]
+            for arr_dtype in all_dtypes]
+
+        for arr_list in all_test_arrays:
+            for arr in arr_list:
+                for axis in (0, 1, 2):
+                    if axis > len(arr.shape)-1:
+                        continue
+                    with self.subTest("Testing np.sum(axis) with {} "
+                                      "input ".format(arr.dtype)):
+                        npy_res = pyfunc(arr, axis=axis)
+                        numba_res = cfunc(arr, axis=axis)
+                        if isinstance(numba_res, np.ndarray):
+                            self.assertPreciseEqual(
+                                npy_res.astype(out_dtypes[arr.dtype]),
+                                numba_res.astype(out_dtypes[arr.dtype]))
+                        else:
+                            # the results are scalars
+                            self.assertEqual(npy_res, numba_res)
+
+    def test_sum_dtype_kws(self):
+        pyfunc = array_sum_dtype_kws
+        cfunc = jit(nopython=True)(pyfunc)
+        all_dtypes = [np.float64, np.float32, np.int64, np.int32, np.uint32,
+                      np.uint64, np.complex64]
+        # timedelta test cannot be enabled until issue #4540 is fixed
+        # all_dtypes += [np.timedelta64]
+        all_test_arrays = [
+            [np.ones((7, 6, 5, 4, 3), arr_dtype),
+             np.ones(1, arr_dtype),
+             np.ones((7, 3), arr_dtype) * -5]
+            for arr_dtype in all_dtypes]
+
+        out_dtypes = {np.dtype('float64'): [np.float64],
+                      np.dtype('float32'): [np.float64, np.float32],
+                      np.dtype('int64'): [np.float64, np.int64, np.float32],
+                      np.dtype('int32'): [np.float64, np.int64, np.float32, np.int32],
+                      np.dtype('uint32'): [np.float64, np.int64, np.float32],
+                      np.dtype('uint64'): [np.float64, np.int64],
+                      np.dtype('complex64'): [np.complex64],
+                      np.dtype('timedelta64'): [np.timedelta64]}
+
+        for arr_list in all_test_arrays:
+            for arr in arr_list:
+                for out_dtype in out_dtypes[arr.dtype]:
+                    subtest_str = ("Testing np.sum with {} input and {} output"
+                                   .format(arr.dtype, out_dtype))
+                    with self.subTest(subtest_str):
+                        self.assertPreciseEqual(pyfunc(arr, dtype=out_dtype),
+                                                cfunc(arr, dtype=out_dtype))
+
+    def test_sum_dtype_kws_negative(self):
+        pyfunc = array_sum_dtype_kws
+        cfunc = jit(nopython=True)(pyfunc)
+        dtype = np.float64
+        # OK
+        a = np.ones((7, 6, 5, 4, 3))
+        self.assertFalse(type(pyfunc(a, dtype=np.int32)) == cfunc(a, dtype=dtype))
+
+    def test_sum_axis_dtype_kws(self):
+        pyfunc = array_sum_axis_dtype_kws
+        cfunc = jit(nopython=True)(pyfunc)
+        dtype = np.float64
+        # OK
+        a = np.ones((7, 6, 5, 4, 3))
+        self.assertPreciseEqual(pyfunc(a, axis=1, dtype=dtype),
+                                cfunc(a,  axis=1, dtype=dtype))
+        self.assertPreciseEqual(pyfunc(a, axis=2, dtype=dtype),
+                                cfunc(a, axis=2, dtype=dtype))
+
+    def test_sum_axis_dtype_kws2(self):
+        pyfunc = array_sum_axis_dtype_kws
+        cfunc = jit(nopython=True)(pyfunc)
+        all_dtypes = [np.float64, np.float32, np.int64, np.int32, np.uint32,
+                      np.uint64, np.complex64]
+        all_test_arrays = [
+            [np.ones((7, 6, 5, 4, 3), arr_dtype),
+             np.ones(1, arr_dtype),
+             np.ones((7, 3), arr_dtype) * -5]
+            for arr_dtype in all_dtypes]
+
+        out_dtypes = {np.dtype('float64'): [np.float64],
+                      np.dtype('float32'): [np.float64, np.float32],
+                      np.dtype('int64'): [np.float64, np.int64, np.float32],
+                      np.dtype('int32'): [np.float64, np.int64, np.float32, np.int32],
+                      np.dtype('uint32'): [np.float64, np.int64, np.float32],
+                      np.dtype('uint64'): [np.float64, np.uint64],
+                      np.dtype('complex64'): [np.complex64],
+                      np.dtype('timedelta64'): [np.timedelta64]}
+
+        for arr_list in all_test_arrays:
+            for arr in arr_list:
+                for out_dtype in out_dtypes[arr.dtype]:
+                    for axis in (0, 1, 2):
+                        if axis > len(arr.shape) - 1:
+                            continue
+                        subtest_str = ("Testing np.sum with {} input and {} output "
+                                       .format(arr.dtype, out_dtype))
+                        with self.subTest(subtest_str):
+                            py_res = pyfunc(arr, axis=axis, dtype=out_dtype)
+                            nb_res = cfunc(arr, axis=axis, dtype=out_dtype)
+                            self.assertPreciseEqual(py_res, nb_res)
+
+    def test_sum_axis_dtype_pos_arg(self):
+        """
+        testing that axis and dtype inputs work when passed as positional
+        :return:
+        """
+        pyfunc = array_sum_axis_dtype_pos
+        cfunc = jit(nopython=True)(pyfunc)
+        dtype = np.float64
+        # OK
+        a = np.ones((7, 6, 5, 4, 3))
+        self.assertPreciseEqual(pyfunc(a, 1, dtype),
+                                cfunc(a,  1, dtype))
+
+        self.assertPreciseEqual(pyfunc(a, 2, dtype),
+                                cfunc(a, 2, dtype))
+
     def test_sum_1d_kws(self):
         # check 1d reduces to scalar
-        pyfunc = array_sum_kws
+        pyfunc = array_sum_axis_kws
         cfunc = jit(nopython=True)(pyfunc)
         a = np.arange(10.)
         self.assertPreciseEqual(pyfunc(a, axis=0), cfunc(a, axis=0))
@@ -894,7 +1079,7 @@ def check(arr, ind):
         test_indices.append(np.array([1, 5, 1, 11, 3]))
         test_indices.append(np.array([[1, 5, 1], [11, 3, 0]], order='F'))
         test_indices.append(np.array([[[1, 5, 1], [11, 3, 0]]]))
-        test_indices.append(np.array([[[[1, 5]], [[11, 0]],[[1, 2]]]]))
+        test_indices.append(np.array([[[[1, 5]], [[11, 0]], [[1, 2]]]]))
         test_indices.append([1, 5, 1, 11, 3])
         test_indices.append((1, 5, 1))
         test_indices.append(((1, 5, 1), (11, 3, 2)))
@@ -914,7 +1099,7 @@ def check(arr, ind):
                            [szA], [-szA - 1]]
         for x in illegal_indices:
             with self.assertRaises(IndexError):
-                cfunc(A, x) # oob raises
+                cfunc(A, x) #  oob raises
 
         # check float indexing raises
         with self.assertRaises(TypingError):
diff --git a/numba/types/functions.py b/numba/types/functions.py
index 1f2c413e9cb..3e91b495da5 100644
--- a/numba/types/functions.py
+++ b/numba/types/functions.py
@@ -79,6 +79,7 @@ class BaseFunction(Callable):
     """
 
     def __init__(self, template):
+
         if isinstance(template, (list, tuple)):
             self.templates = tuple(template)
             keys = set(temp.key for temp in self.templates)
diff --git a/numba/typing/arraydecl.py b/numba/typing/arraydecl.py
index 4a684f6b260..fd0a60d5076 100644
--- a/numba/typing/arraydecl.py
+++ b/numba/typing/arraydecl.py
@@ -612,22 +612,38 @@ def sum_expand(self, args, kws):
     sum can be called with or without an axis parameter.
     """
     pysig = None
-    if kws:
+    if 'axis' in kws and not 'dtype' in kws:
         def sum_stub(axis):
             pass
         pysig = utils.pysignature(sum_stub)
         # rewrite args
         args = list(args) + [kws['axis']]
-        kws = None
+        # kws = None
+
+    if 'dtype' in kws and not 'axis' in kws:
+        def sum_stub(dtype):
+            pass
+        pysig = utils.pysignature(sum_stub)
+        # rewrite args
+        args = list(args) + [kws['dtype']]
+        # kws = None
+
+    if 'dtype' in kws and 'axis' in kws:
+        def sum_stub(axis, dtype):
+            pass
+        pysig = utils.pysignature(sum_stub)
+        # rewrite args
+        args = list(args) + [kws['axis'], kws['dtype']]
+
     args_len = len(args)
-    assert args_len <= 1
+    assert args_len <= 2
     if args_len == 0:
-        # No axis parameter so the return type of the summation is a scalar
+        # No axis or dtype parameter so the return type of the summation is a scalar
         # of the type of the array.
         out = signature(_expand_integer(self.this.dtype), *args,
                         recvr=self.this)
-    else:
-        # There is an axis parameter
+    elif args_len == 1 and not 'dtype' in kws:
+        # There is an axis parameter, either arg or kwarg
         if self.this.ndim == 1:
             # 1d reduces to a scalar
             return_type = self.this.dtype
@@ -637,6 +653,30 @@ def sum_stub(axis):
             return_type = types.Array(dtype=_expand_integer(self.this.dtype),
                                     ndim=self.this.ndim-1, layout='C')
         out = signature(return_type, *args, recvr=self.this)
+
+    elif args_len == 1 and 'dtype' in kws:
+        # No axis parameter so the return type of the summation is a scalar
+        # of the dtype parameter.
+        from .npydecl import _parse_dtype
+        dtype, = args
+        dtype = _parse_dtype(dtype)
+        out = signature(dtype, *args, recvr=self.this)
+
+    elif args_len == 2:
+        # There is an axis and dtype parameter, either arg or kwarg
+        from .npydecl import _parse_dtype
+        dtype = _parse_dtype(args[1])
+        return_type = dtype
+        if self.this.ndim != 1:
+            # 1d reduces to a scalar, 2d and above reduce dim by 1
+            # the return type of this summation is  an array of dimension one
+            # less than the input array.
+            return_type = types.Array(dtype=return_type,
+                                    ndim=self.this.ndim-1, layout='C')
+        out = signature(return_type, *args, recvr=self.this)
+
+    else:
+        pass
     return out.replace(pysig=pysig)
 
 def generic_expand_cumulative(self, args, kws):
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
index bde8225600c..77292c40785 100644
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -378,9 +378,18 @@ def generic(self, args, kws):
         pysig = None
         if kws:
             if self.method_name == 'sum':
-                def sum_stub(arr, axis):
-                    pass
-                pysig = utils.pysignature(sum_stub)
+                if 'axis' in kws and not 'dtype' in kws:
+                    def sum_stub(arr, axis):
+                        pass
+                    pysig = utils.pysignature(sum_stub)
+                if 'dtype' in kws and not 'axis' in kws:
+                    def sum_stub(arr, dtype):
+                        pass
+                    pysig = utils.pysignature(sum_stub)
+                if 'dtype' in kws and 'axis' in kws:
+                    def sum_stub(arr, axis, dtype):
+                        pass
+                    pysig = utils.pysignature(sum_stub)
             elif self.method_name == 'argsort':
                 def argsort_stub(arr, kind='quicksort'):
                     pass


Code review comments:
diff_hunk: @@ -378,9 +378,18 @@ def generic(self, args, kws):
         pysig = None
         if kws:
             if self.method_name == 'sum':
-                def sum_stub(arr, axis):
-                    pass
-                pysig = utils.pysignature(sum_stub)
+                if 'axis' in kws and not 'dtype' in kws:
path: numba/typing/npydecl.py
text: @user1:
I think the preferred syntax is `'dtype' not in kws`, it is also slightly more readable.

@author:
you're right, changed.

diff_hunk: @@ -378,9 +378,18 @@ def generic(self, args, kws):
         pysig = None
         if kws:
             if self.method_name == 'sum':
-                def sum_stub(arr, axis):
-                    pass
-                pysig = utils.pysignature(sum_stub)
+                if 'axis' in kws and not 'dtype' in kws:
+                    def sum_stub(arr, axis):
+                        pass
+                    pysig = utils.pysignature(sum_stub)
+                if 'dtype' in kws and not 'axis' in kws:
+                    def sum_stub(arr, dtype):
+                        pass
+                    pysig = utils.pysignature(sum_stub)
+                if 'dtype' in kws and 'axis' in kws:
path: numba/typing/npydecl.py
text: @user1:
This could also be an `elif` clause.

@author:
yes, fixed

diff_hunk: @@ -612,22 +612,38 @@ def sum_expand(self, args, kws):
     sum can be called with or without an axis parameter.
     """
     pysig = None
-    if kws:
+    if 'axis' in kws and not 'dtype' in kws:
path: numba/typing/arraydecl.py
text: @user1:
Prefer to use `not in`.

@author:
done

diff_hunk: @@ -612,22 +612,38 @@ def sum_expand(self, args, kws):
     sum can be called with or without an axis parameter.
     """
     pysig = None
-    if kws:
+    if 'axis' in kws and not 'dtype' in kws:
         def sum_stub(axis):
             pass
         pysig = utils.pysignature(sum_stub)
         # rewrite args
         args = list(args) + [kws['axis']]
-        kws = None
+        # kws = None
path: numba/typing/arraydecl.py
text: @user1:
dead code?

@author:
tbh, I couldn't understand why that was there so left it in case I had to restore it. Since everything is working I can delete it. However, it's that feeling you get when you put together a machine and there's a piece left outside. I'll delete it.

diff_hunk: @@ -612,22 +612,38 @@ def sum_expand(self, args, kws):
     sum can be called with or without an axis parameter.
     """
     pysig = None
-    if kws:
+    if 'axis' in kws and not 'dtype' in kws:
         def sum_stub(axis):
             pass
         pysig = utils.pysignature(sum_stub)
         # rewrite args
         args = list(args) + [kws['axis']]
-        kws = None
+        # kws = None
+
+    if 'dtype' in kws and not 'axis' in kws:
path: numba/typing/arraydecl.py
text: @user1:
Prefer to use `not in` and candidate for `elif`.

@author:
done

diff_hunk: @@ -612,22 +612,38 @@ def sum_expand(self, args, kws):
     sum can be called with or without an axis parameter.
     """
     pysig = None
-    if kws:
+    if 'axis' in kws and not 'dtype' in kws:
         def sum_stub(axis):
             pass
         pysig = utils.pysignature(sum_stub)
         # rewrite args
         args = list(args) + [kws['axis']]
-        kws = None
+        # kws = None
+
+    if 'dtype' in kws and not 'axis' in kws:
+        def sum_stub(dtype):
+            pass
+        pysig = utils.pysignature(sum_stub)
+        # rewrite args
+        args = list(args) + [kws['dtype']]
+        # kws = None
path: numba/typing/arraydecl.py
text: @user1:
Dead code?

@author:
deleted

diff_hunk: @@ -612,22 +612,38 @@ def sum_expand(self, args, kws):
     sum can be called with or without an axis parameter.
     """
     pysig = None
-    if kws:
+    if 'axis' in kws and not 'dtype' in kws:
         def sum_stub(axis):
             pass
         pysig = utils.pysignature(sum_stub)
         # rewrite args
         args = list(args) + [kws['axis']]
-        kws = None
+        # kws = None
+
+    if 'dtype' in kws and not 'axis' in kws:
+        def sum_stub(dtype):
+            pass
+        pysig = utils.pysignature(sum_stub)
+        # rewrite args
+        args = list(args) + [kws['dtype']]
+        # kws = None
+
+    if 'dtype' in kws and 'axis' in kws:
path: numba/typing/arraydecl.py
text: @user1:
Candidate for `elif`.

diff_hunk: @@ -612,22 +612,38 @@ def sum_expand(self, args, kws):
     sum can be called with or without an axis parameter.
path: numba/typing/arraydecl.py
text: @user1:
This documentation line should probably be updated to reflect the new signature.

@author:
fixed

diff_hunk: @@ -268,8 +226,129 @@ def array_sum_impl_axis(arr, axis):
                 elif axis == 3:
                     index_tuple4 = _gen_index_tuple(arr.shape, axis_index, 3)
                     result += arr[index_tuple4]
-
         return op(result, 0)
+    return inner
+
+
+@lower_builtin(np.sum, types.Array, types.intp, types.DTypeSpec)
+@lower_builtin(np.sum, types.Array, types.IntegerLiteral, types.DTypeSpec)
+@lower_builtin("array.sum", types.Array, types.intp, types.DTypeSpec)
+@lower_builtin("array.sum", types.Array, types.IntegerLiteral, types.DTypeSpec)
+def array_sum_axis_dtype(context, builder, sig, args):
+    """
+    The third parameter to gen_index_tuple that generates the indexing
+    tuples has to be a const so we can't just pass "axis" through since
+    that isn't const.  We can check for specific values and have
+    different instances that do take consts.  Supporting axis summation
+    only up to the fourth dimension for now.
+    """
+    # typing/arraydecl.py:sum_expand defines the return type for sum with axis.
+    # It is one dimension less than the input array.
+
+    retty = sig.return_type
+    zero = getattr(retty, 'dtype', retty)(0)
+    # if the return is scalar in type then "take" the 0th element of the
+    # 0d array accumulator as the return value
+    if getattr(retty, 'ndim', None) is None:
+        op = np.take
+    else:
+        op = _array_sum_axis_nop
+    [ty_array, ty_axis, ty_dtype] = sig.args
+    is_axis_const = False
+    const_axis_val = 0
+    if isinstance(ty_axis, types.Literal):
+        # this special-cases for constant axis
+        const_axis_val = ty_axis.literal_value
+        # fix negative axis
+        if const_axis_val < 0:
+            const_axis_val = ty_array.ndim + const_axis_val
+        if const_axis_val < 0 or const_axis_val > ty_array.ndim:
+            raise ValueError("'axis' entry is out of bounds")
+
+        ty_axis = context.typing_context.resolve_value_type(const_axis_val)
+        axis_val = context.get_constant(ty_axis, const_axis_val)
+        # rewrite arguments
+        args = args[0], axis_val, args[2]
+        # rewrite sig
+        sig = sig.replace(args=[ty_array, ty_axis, ty_dtype])
+        is_axis_const = True
+
+    gen_impl = gen_sum_axis_impl(is_axis_const, const_axis_val, op, zero)
+    compiled = register_jitable(gen_impl)
+
+    def array_sum_impl_axis(arr, axis, dtype):
+        return compiled(arr, axis)
+
+    res = context.compile_internal(builder, array_sum_impl_axis, sig, args)
+    return impl_ret_new_ref(context, builder, sig.return_type, res)
+
+
+@lower_builtin(np.sum, types.Array,  types.DTypeSpec)
+@lower_builtin(np.sum, types.Array,  types.DTypeSpec)
+@lower_builtin("array.sum", types.Array, types.DTypeSpec)
path: numba/targets/arraymath.py
text: @user1:
This also seems duplicated

@author:
fixed

diff_hunk: @@ -378,9 +378,18 @@ def generic(self, args, kws):
         pysig = None
         if kws:
             if self.method_name == 'sum':
-                def sum_stub(arr, axis):
-                    pass
-                pysig = utils.pysignature(sum_stub)
+                if 'axis' in kws and not 'dtype' in kws:
+                    def sum_stub(arr, axis):
+                        pass
+                    pysig = utils.pysignature(sum_stub)
+                if 'dtype' in kws and not 'axis' in kws:
path: numba/typing/npydecl.py
text: @user1:
Same comment as above on the negation, also, this could be an `elif` clause.

@author:
fixed