Repository: cupy/cupy
Pull request: Fix dtype option of sum and prod
This PR fixes the logic of `dtype` argument of `sum` and `prod`, so that it enables e.g.
`cupy.sum(cupy.arange(3).astype(cupy.int16), dtype=cupy.int32)`

close #1240.

Problem statement:
sum without upcast
`cupy.sum` and `cupy.prod` upcasts ints (or bool) to `int64` or `uint64`, to align with numpy.  This feature would be disabled with `x.sum(dtype=x.dtype)` but not supported in cupy.
```
>>> x = cupy.arange(3).astype(cupy.int16)
>>> x.sum(dtype=x.dtype)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "cupy/core/core.pyx", line 1139, in cupy.core.core.ndarray.sum
  File "cupy/core/core.pyx", line 1147, in cupy.core.core.ndarray.sum
  File "cupy/core/reduction.pxi", line 222, in cupy.core.core.simple_reduction_function.__call__
  File "cupy/core/elementwise.pxi", line 698, in cupy.core.core._guess_routine
TypeError: Wrong type (<class 'numpy.int16'>) of arguments for cupy_sum
>>> cupy.__version__
'4.0.0'
```

Hints from issue discussion:

Commit message: remove tests between real and compex

Patch to review:
diff --git a/cupy/core/core.pyx b/cupy/core/core.pyx
index 03f901e72d4..183d087f397 100644
--- a/cupy/core/core.pyx
+++ b/cupy/core/core.pyx
@@ -1144,7 +1144,10 @@ cdef class ndarray:
            :meth:`numpy.ndarray.sum`
 
         """
-        return _sum(self, axis, dtype, out, keepdims)
+        if dtype is None:
+            return _sum_auto_dtype(self, axis, dtype, out, keepdims)
+        else:
+            return _sum_keep_dtype(self, axis, dtype, out, keepdims)
 
     # TODO(okuta): Implement cumsum
 
@@ -1190,7 +1193,10 @@ cdef class ndarray:
            :meth:`numpy.ndarray.prod`
 
         """
-        return _prod(self, axis, dtype, out, keepdims)
+        if dtype is None:
+            return _prod_auto_dtype(self, axis, dtype, out, keepdims)
+        else:
+            return _prod_keep_dtype(self, axis, dtype, out, keepdims)
 
     # TODO(okuta): Implement cumprod
 
@@ -3341,7 +3347,7 @@ cpdef _prepare_multiple_array_indexing(ndarray a, list slices):
         [index._reshape((1,) + index.shape) for index in flattened_indexes],
         axis=0, shape=concat_shape, dtype=flattened_indexes[0].dtype)
 
-    reduced_idx = _sum(flattened_indexes, axis=0)
+    reduced_idx = _sum_auto_dtype(flattened_indexes, axis=0)
 
     return a_interm, reduced_idx, li, ri
 
@@ -3966,21 +3972,39 @@ _any = create_reduction_func(
 # Mathematical functions
 # -----------------------------------------------------------------------------
 
-_sum = create_reduction_func(
+_sum_auto_dtype = create_reduction_func(
     'cupy_sum',
-    ('?->l', 'B->L', 'h->l', 'H->L', 'i->l', 'I->L', 'l->l', 'L->L',
+    ('?->l', 'b->l', 'B->L', 'h->l', 'H->L', 'i->l', 'I->L', 'l->l', 'L->L',
      'q->q', 'Q->Q',
      ('e->e', (None, None, None, 'float')),
      'f->f', 'd->d', 'F->F', 'D->D'),
     ('in0', 'a + b', 'out0 = type_out0_raw(a)', None), 0)
 
 
-_prod = create_reduction_func(
+_sum_keep_dtype = create_reduction_func(
+    'cupy_sum',
+    ('?->?', 'b->b', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L',
+     'q->q', 'Q->Q',
+     ('e->e', (None, None, None, 'float')),
+     'f->f', 'd->d', 'F->F', 'D->D'),
+    ('in0', 'a + b', 'out0 = type_out0_raw(a)', None), 0)
+
+
+_prod_auto_dtype = create_reduction_func(
+    'cupy_prod',
+    ('?->l', 'b->l', 'B->L', 'h->l', 'H->L', 'i->l', 'I->L', 'l->l', 'L->L',
+     'q->q', 'Q->Q',
+     ('e->e', (None, None, None, 'float')),
+     'f->f', 'd->d', 'F->F', 'D->D'),
+    ('in0', 'a * b', 'out0 = type_out0_raw(a)', None), 1)
+
+
+_prod_keep_dtype = create_reduction_func(
     'cupy_prod',
-    ['?->l', 'B->L', 'h->l', 'H->L', 'i->l', 'I->L', 'l->l', 'L->L',
+    ('?->?', 'b->b', 'B->B', 'h->h', 'H->H', 'i->i', 'I->I', 'l->l', 'L->L',
      'q->q', 'Q->Q',
      ('e->e', (None, None, None, 'float')),
-     'f->f', 'd->d', 'F->F', 'D->D'],
+     'f->f', 'd->d', 'F->F', 'D->D'),
     ('in0', 'a * b', 'out0 = type_out0_raw(a)', None), 1)
 
 
diff --git a/cupy/core/fusion.py b/cupy/core/fusion.py
index 8815c5a2bb4..07a3dd1c2a0 100644
--- a/cupy/core/fusion.py
+++ b/cupy/core/fusion.py
@@ -851,7 +851,7 @@ def __call__(self, *args, **kwargs):
 
 all._raw = core._all
 any._raw = core._any
-sum._raw = core._sum
-prod._raw = core._prod
+sum._raw = core._sum_auto_dtype
+prod._raw = core._prod_auto_dtype
 amax._raw = core._amax
 amin._raw = core._amin
diff --git a/tests/cupy_tests/math_tests/test_sumprod.py b/tests/cupy_tests/math_tests/test_sumprod.py
index b98521dd344..ab78ccd5d74 100644
--- a/tests/cupy_tests/math_tests/test_sumprod.py
+++ b/tests/cupy_tests/math_tests/test_sumprod.py
@@ -118,6 +118,16 @@ def test_sum_axes4(self, xp, dtype):
         a = testing.shaped_arange((20, 30, 40, 50), xp, dtype)
         return a.sum(axis=(0, 2, 3))
 
+    @testing.for_all_dtypes_combination(names=['src_dtype', 'dst_dtype'])
+    @testing.numpy_cupy_allclose()
+    def test_sum_dtype(self, xp, src_dtype, dst_dtype):
+        if not xp.can_cast(src_dtype, dst_dtype) \
+                or [xp.dtype(src_dtype).kind,
+                    xp.dtype(dst_dtype).kind].count('c') == 1:
+            return xp.array([])  # skip
+        a = testing.shaped_arange((2, 3, 4), xp, src_dtype)
+        return a.sum(dtype=dst_dtype)
+
     @testing.numpy_cupy_allclose()
     def test_sum_keepdims(self, xp):
         a = testing.shaped_arange((2, 3, 4), xp)
@@ -161,6 +171,16 @@ def test_external_prod_axis(self, xp, dtype):
         a = testing.shaped_arange((2, 3, 4), xp, dtype)
         return xp.prod(a, axis=1)
 
+    @testing.for_all_dtypes_combination(names=['src_dtype', 'dst_dtype'])
+    @testing.numpy_cupy_allclose()
+    def test_prod_dtype(self, xp, src_dtype, dst_dtype):
+        if not xp.can_cast(src_dtype, dst_dtype) \
+                or [xp.dtype(src_dtype).kind,
+                    xp.dtype(dst_dtype).kind].count('c') == 1:
+            return xp.array([])  # skip
+        a = testing.shaped_arange((2, 3), xp, src_dtype)
+        return a.prod(dtype=dst_dtype)
+
 
 axes = [0, 1, 2]
 


Code review comments:
diff_hunk: @@ -118,6 +118,16 @@ def test_sum_axes4(self, xp, dtype):
         a = testing.shaped_arange((20, 30, 40, 50), xp, dtype)
         return a.sum(axis=(0, 2, 3))
 
+    @testing.for_all_dtypes_combination(names=['src_dtype', 'dst_dtype'])
+    @testing.numpy_cupy_allclose()
+    def test_sum_dtype(self, xp, src_dtype, dst_dtype):
+        if not xp.can_cast(src_dtype, dst_dtype) \
+                or [xp.dtype(src_dtype).kind,
+                    xp.dtype(dst_dtype).kind].count('c') == 1:
path: tests/cupy_tests/math_tests/test_sumprod.py
text: @user1:
I guess it should be `src_dtype == xp.float16 and xp.dtype(dst_dtype).kind == 'c'`?