Repository: cupy/cupy
Pull request: keepdims should always preserve all dimensions in CUB-based reductions
closes #2720

The bug is fixed by the lines computing `out_shape` within `device_reduce`.

In order to avoid multiple calls to `_get_axis`, I moved the `_get_axis` call outside of `_preprocess_array` and call it directly in `_routines_math.pyx` and `_routines_statistics.pyx` instead. The duplicate logic could potentially be refactored into a common function

`_get_axes` will have already sorted the axes in ascending order and checked their range (It also separates out the axes to be reduced from those needed in determining the shape for keepdims).
This allows greatly simplifying the code in `_cub_device_segmented_reduce_axis_compatible` and removes the need for `_cub_device_reduce_axis_compatible` altogether.

Although not strictly part of the bug fix, I implemented a performance improvement in the segmented reduction case by replacing the following slow call

```Python
numpy.all(numpy.diff(cub_axis) == 1)
```
with a call to a new `_contig_axis` Cython helper function


Problem statement:
inconsistent behavior of keepdims for CUB-based reductions
It appears that keepdims does not keep all of the reduction dimensions when CUB is enabled and axis is a tuple containing more than one axis.

* Conditions (you can just paste the output of `python -c 'import cupy; cupy.show_config()'`)
```
CuPy Version          :  7.0.0rc1 (build from master branch)
CUDA Root             : /usr/local/cuda
CUDA Build Version    : 10010
CUDA Driver Version   : 10010
CUDA Runtime Version  : 10010

```
* Code to reproduce

```Python
import cupy
x = cupy.ones((64, 64))

cupy.cuda.cub_enabled=False
out = x.sum(axis=(0, 1), keepdims=True)
print("CUB DISABLED shape={}".format(out.shape))

cupy.cuda.cub_enabled=True
out = x.sum(axis=(0, 1), keepdims=True)
print("CUB ENABLED shape={}".format(out.shape))
```
gives
```
CUB DISABLED shape=(1, 1)
CUB ENABLED shape=(1,)
```

cc @leofang

Hints from issue discussion:
actually the behavior is as above only for `sum`, `min` and `max`

For `argmin` and `argmax` the CUB case gives shape `()` instead( while non-CUB is still (1, 1) as expected)
The issue is within `device_reduce` rather than `device_segmented_reduce`. 

This can be confirmed by using `axis=(1, 0)` instead of `axis=(0,1)` which does not meet the criteria of `_cub_device_reduce_axis_compatible` . In that case, all reductions give shape (1, 1) as expected indicating that the `device_segmented_reduce` treats `keepdims` correctly.


A separate issue is to fix `_cub_device_reduce_axis_compatible` so that axis (0, 1) and (1,0) are both recognized as compatible.
I am working on a fix for this currently.
Apparently I didn’t think it was a problem 🤦‍♂️ Thanks for catching it, @grlee77. I think calling `_preprocess_array()` should give you the right shape? 
> Apparently I didn’t think it was a problem

I think the origin of the problem was an oversight on my part in #2517. I enabled reduction over all axes there, but did not ensure that `keepdims` preserved all of the axes.

Commit message: remove unused import

Patch to review:
diff --git a/cupy/core/_reduction.pyx b/cupy/core/_reduction.pyx
index 7933d1cfc01..2f2420e6f80 100644
--- a/cupy/core/_reduction.pyx
+++ b/cupy/core/_reduction.pyx
@@ -109,7 +109,7 @@ extern "C" __global__ void ${name}(${params}) {
 cpdef tuple _get_axis(object axis, Py_ssize_t ndim):
     cdef Py_ssize_t dim
     if axis is None:
-        axis = tuple(range(ndim))
+        return (tuple(range(ndim)), ())
     elif sequence.PySequence_Check(axis):
         axis = tuple(axis)
     else:
diff --git a/cupy/core/_routines_math.pyx b/cupy/core/_routines_math.pyx
index c35015d0f01..ffe12754f1d 100644
--- a/cupy/core/_routines_math.pyx
+++ b/cupy/core/_routines_math.pyx
@@ -74,21 +74,11 @@ cdef ndarray _ndarray_prod(ndarray self, axis, dtype, out, keepdims):
 
 cdef ndarray _ndarray_sum(ndarray self, axis, dtype, out, keepdims):
     if cupy.cuda.cub_enabled:
-        if cub.can_use_device_reduce(cub.CUPY_CUB_SUM, self.dtype, self.ndim,
-                                     axis, dtype):
-            return cub.device_reduce(self, cub.CUPY_CUB_SUM, out=out,
-                                     keepdims=keepdims)
-
-        if self.flags.c_contiguous:
-            order = 'C'
-        elif self.flags.f_contiguous:
-            order = 'F'
-        else:
-            order = None
-        if cub.can_use_device_segmented_reduce(cub.CUPY_CUB_SUM, self.dtype,
-                                               self.ndim, axis, dtype, order):
-            return cub.device_segmented_reduce(self, cub.CUPY_CUB_SUM, axis,
-                                               out=out, keepdims=keepdims)
+        # result will be None if the reduction is not compatible with CUB
+        result = cub.cub_reduction(self, cub.CUPY_CUB_SUM, axis, dtype, out,
+                                   keepdims)
+        if result is not None:
+            return result
     if dtype is None:
         return _sum_auto_dtype(self, axis, dtype, out, keepdims)
     else:
diff --git a/cupy/core/_routines_statistics.pyx b/cupy/core/_routines_statistics.pyx
index 1fc6d385df9..6ebd5b070c9 100644
--- a/cupy/core/_routines_statistics.pyx
+++ b/cupy/core/_routines_statistics.pyx
@@ -14,68 +14,43 @@ if cupy.cuda.cub_enabled:
 
 cdef ndarray _ndarray_max(ndarray self, axis, out, dtype, keepdims):
     if cupy.cuda.cub_enabled:
-        if cub.can_use_device_reduce(cub.CUPY_CUB_MAX, self.dtype, self.ndim,
-                                     axis, dtype):
-            return cub.device_reduce(self, cub.CUPY_CUB_MAX, out=out,
-                                     keepdims=keepdims)
-
-        if self.flags.c_contiguous:
-            order = 'C'
-        elif self.flags.f_contiguous:
-            order = 'F'
-        else:
-            order = None
-        if cub.can_use_device_segmented_reduce(cub.CUPY_CUB_MAX, self.dtype,
-                                               self.ndim, axis, dtype, order):
-            return cub.device_segmented_reduce(self, cub.CUPY_CUB_MAX, axis,
-                                               out=out, keepdims=keepdims)
+        # result will be None if the reduction is not compatible with CUB
+        result = cub.cub_reduction(self, cub.CUPY_CUB_MAX, axis, dtype, out,
+                                   keepdims)
+        if result is not None:
+            return result
     return _amax(self, axis=axis, out=out, dtype=dtype, keepdims=keepdims)
 
 
 cdef ndarray _ndarray_min(ndarray self, axis, out, dtype, keepdims):
     if cupy.cuda.cub_enabled:
-        if cub.can_use_device_reduce(cub.CUPY_CUB_MIN, self.dtype, self.ndim,
-                                     axis, dtype):
-            return cub.device_reduce(self, cub.CUPY_CUB_MIN, out=out,
-                                     keepdims=keepdims)
-        if self.flags.c_contiguous:
-            order = 'C'
-        elif self.flags.f_contiguous:
-            order = 'F'
-        else:
-            order = None
-        if cub.can_use_device_segmented_reduce(cub.CUPY_CUB_MIN, self.dtype,
-                                               self.ndim, axis, dtype, order):
-            return cub.device_segmented_reduce(self, cub.CUPY_CUB_MIN, axis,
-                                               out=out, keepdims=keepdims)
+        # result will be None if the reduction is not compatible with CUB
+        result = cub.cub_reduction(self, cub.CUPY_CUB_MIN, axis, out, dtype,
+                                   keepdims)
+        if result is not None:
+            return result
     return _amin(self, axis=axis, out=out, dtype=dtype, keepdims=keepdims)
 
 
 # TODO(leofang): this signature is incompatible with NumPy!
 cdef ndarray _ndarray_argmax(ndarray self, axis, out, dtype, keepdims):
     if cupy.cuda.cub_enabled:
-        # Note that the NumPy signature of argmax only has axis and out, so we
-        # need to disable the rest. Moreover, to be compatible with NumPy, axis
-        # can only be None or integers
-        if cub.can_use_device_reduce(
-                cub.CUPY_CUB_ARGMAX, self.dtype, self.ndim, axis, None):
-            return cub.device_reduce(self, cub.CUPY_CUB_ARGMAX, out=out,
-                                     keepdims=False)
-        # TODO(leofang): support device_segmented_reduce for axis=-1?
+        # result will be None if the reduction is not compatible with CUB
+        result = cub.cub_reduction(self, cub.CUPY_CUB_ARGMAX, axis, dtype, out,
+                                   keepdims)
+        if result is not None:
+            return result
     return _argmax(self, axis=axis, out=out, dtype=dtype, keepdims=keepdims)
 
 
 # TODO(leofang): this signature is incompatible with NumPy!
 cdef ndarray _ndarray_argmin(ndarray self, axis, out, dtype, keepdims):
     if cupy.cuda.cub_enabled:
-        # Note that the NumPy signature of argmax only has axis and out, so we
-        # need to disable the rest. Moreover, to be compatible with NumPy, axis
-        # can only be None or integers
-        if cub.can_use_device_reduce(
-                cub.CUPY_CUB_ARGMIN, self.dtype, self.ndim, axis, None):
-            return cub.device_reduce(self, cub.CUPY_CUB_ARGMIN, out=out,
-                                     keepdims=False)
-        # TODO(leofang): support device_segmented_reduce for axis=-1?
+        # result will be None if the reduction is not compatible with CUB
+        result = cub.cub_reduction(self, cub.CUPY_CUB_ARGMIN, axis, dtype, out,
+                                   keepdims)
+        if result is not None:
+            return result
     return _argmin(self, axis=axis, out=out, dtype=dtype, keepdims=keepdims)
 
 
diff --git a/cupy/cuda/cub.pyx b/cupy/cuda/cub.pyx
index 330ba9a07f3..835bc4c765b 100644
--- a/cupy/cuda/cub.pyx
+++ b/cupy/cuda/cub.pyx
@@ -2,6 +2,8 @@
 
 """Wrapper of CUB functions for CuPy API."""
 
+from cpython cimport sequence
+
 import numpy
 
 from cupy.core.core cimport ndarray, _internal_ascontiguousarray
@@ -61,18 +63,40 @@ cdef extern from 'cupy_cub.h' nogil:
 # Python interface
 ###############################################################################
 
-cpdef _preprocess_array(ndarray arr, axis, bint keepdims, str order):
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cpdef bint _contig_axes(tuple axes):
+    # True if the specified axes are in ascending order without gaps
+    cdef Py_ssize_t n
+    cdef bint contig = True
+    for n in range(1, len(axes)):
+        contig = (axes[n] - axes[n - 1]) == 1
+        if not contig:
+            break
+    return contig
+
+
+cdef tuple _get_output_shape(ndarray arr, tuple out_axis, bint keepdims):
+    cdef tuple out_shape
+
+    if not keepdims:
+        out_shape = tuple([arr.shape[axis] for axis in out_axis])
+    else:
+        out_shape = tuple([arr.shape[axis] if axis in out_axis else 1
+                           for axis in range(arr.ndim)])
+    return out_shape
+
+
+cpdef Py_ssize_t _preprocess_array(ndarray arr, tuple reduce_axis,
+                                   tuple out_axis, str order):
     '''
     This function more or less follows the logic of _get_permuted_args() in
     reduction.pxi. The input array arr is C- or F- contiguous along axis.
     '''
-    # if import at the top level, a segfault would happen when import cupy!
-    from cupy.core._kernel import _get_axis
-
-    cdef tuple reduce_axis, out_axis, axis_permutes, out_shape
+    cdef tuple axis_permutes, out_shape
     cdef Py_ssize_t contiguous_size = 1
 
-    reduce_axis, out_axis = _get_axis(axis, arr._shape.size())
     # one more sanity check?
     if order == 'C':
         axis_permutes = out_axis + reduce_axis
@@ -82,16 +106,11 @@ cpdef _preprocess_array(ndarray arr, axis, bint keepdims, str order):
 
     for axis in reduce_axis:
         contiguous_size *= arr.shape[axis]
-    if not keepdims:
-        out_shape = tuple([arr.shape[axis] for axis in out_axis])
-    else:
-        out_shape = tuple([arr.shape[axis] if axis in out_axis else 1
-                           for axis in range(arr.ndim)])
-
-    return out_shape, contiguous_size
+    return contiguous_size
 
 
-def device_reduce(ndarray x, int op, out=None, bint keepdims=False):
+def device_reduce(ndarray x, int op, tuple out_axis, out=None,
+                  bint keepdims=False):
     cdef ndarray y, z
     cdef memory.MemoryPointer ws
     cdef int dtype_id, ndim_out, kv_bytes, x_size
@@ -100,8 +119,14 @@ def device_reduce(ndarray x, int op, out=None, bint keepdims=False):
     cdef void *y_ptr
     cdef void *ws_ptr
     cdef Stream_t s
+    cdef tuple out_shape
+
+    if keepdims:
+        out_shape = _get_output_shape(x, out_axis, keepdims)
+        ndim_out = len(out_shape)
+    else:
+        ndim_out = 0
 
-    ndim_out = keepdims
     if out is not None and out.ndim != ndim_out:
         raise ValueError(
             'output parameter for reduction operation has the wrong number of '
@@ -139,15 +164,15 @@ def device_reduce(ndarray x, int op, out=None, bint keepdims=False):
         y = y.reshape(())
 
     if keepdims:
-        y = y.reshape((1,))
+        y = y.reshape(out_shape)
     if out is not None:
         out[...] = y
         y = out
     return y
 
 
-def device_segmented_reduce(ndarray x, int op, axis, out=None,
-                            bint keepdims=False):
+def device_segmented_reduce(ndarray x, int op, tuple reduce_axis,
+                            tuple out_axis, out=None, bint keepdims=False):
     # if import at the top level, a segfault would happen when import cupy!
     from cupy.creation.ranges import arange
 
@@ -178,7 +203,8 @@ def device_segmented_reduce(ndarray x, int op, axis, out=None,
         raise RuntimeError('input is neither C- nor F- contiguous.')
 
     # prepare input
-    out_shape, contiguous_size = _preprocess_array(x, axis, keepdims, order)
+    contiguous_size = _preprocess_array(x, reduce_axis, out_axis, order)
+    out_shape = _get_output_shape(x, out_axis, keepdims)
     x_ptr = <void*>x.data.ptr
     y = ndarray(out_shape, dtype=x.dtype, order=order)
     y_ptr = <void*>y.data.ptr
@@ -269,57 +295,30 @@ def device_csrmv(int n_rows, int n_cols, int nnz, ndarray values,
     return y
 
 
-cdef bint _cub_device_reduce_axis_compatible(axis, Py_ssize_t ndim):
-    if ((axis is None) or ndim == 1 or axis == tuple(range(ndim))):
-        return True
-    return False
-
-
 cdef bint _cub_device_segmented_reduce_axis_compatible(
-        axis, Py_ssize_t ndim, order):
+        tuple cub_axis, Py_ssize_t ndim, order):
     # Implementation borrowed from cupy.fft.fft._get_cufft_plan_nd().
     # This function checks if the reduced axes are contiguous.
-    cdef tuple cub_axis
-
-    if axis is None:
-        # this is impossible, just in case
-        raise ValueError('axis cannot be None.')
-    else:
-        if numpy.isscalar(axis):
-            axis = (axis,)
-        axis = tuple(axis)
-
-        if numpy.min(axis) < -ndim or numpy.max(axis) > ndim - 1:
-            raise ValueError('The specified axis exceed the array dimensions.')
-
-        # sort the provided axis in ascending order
-        cub_axis = tuple(sorted(numpy.mod(axis, ndim)))
-
-        # the axes to be reduced must be C- or F- contiguous
-        if not numpy.all(numpy.diff(cub_axis) == 1):
-            return False
-        if order not in ('c', 'C', 'f', 'F'):
-            return False
-        if order in ('c', 'C') and ((ndim - 1) not in cub_axis):
-            return False
-        if order in ('f', 'F') and (0 not in cub_axis):
-            return False
 
-    return True
+    # the axes to be reduced must be C- or F- contiguous
+    if _contig_axes(cub_axis):
+        if order in ('c', 'C'):
+            return ((ndim - 1) in cub_axis)
+        elif order in ('f', 'F'):
+            return (0 in cub_axis)
+    return False
 
 
-def can_use_device_reduce(int op, x_dtype, Py_ssize_t ndim, axis=None,
-                          dtype=None):
-    if not _cub_reduce_dtype_compatible(x_dtype, op, dtype):
-        return False
-    return _cub_device_reduce_axis_compatible(axis, ndim)
+def can_use_device_reduce(int op, x_dtype, tuple out_axis, dtype=None):
+    return out_axis is () and _cub_reduce_dtype_compatible(x_dtype, op, dtype)
 
 
-def can_use_device_segmented_reduce(int op, x_dtype, Py_ssize_t ndim, axis,
-                                    dtype=None, order='C'):
+def can_use_device_segmented_reduce(int op, x_dtype, Py_ssize_t ndim,
+                                    reduce_axis, dtype=None, order='C'):
     if not _cub_reduce_dtype_compatible(x_dtype, op, dtype):
         return False
-    return _cub_device_segmented_reduce_axis_compatible(axis, ndim, order)
+    return _cub_device_segmented_reduce_axis_compatible(reduce_axis, ndim,
+                                                        order)
 
 
 cdef _cub_reduce_dtype_compatible(x_dtype, int op, dtype=None,
@@ -344,6 +343,61 @@ cdef _cub_reduce_dtype_compatible(x_dtype, int op, dtype=None,
     return True
 
 
+def cub_reduction(arr, op, axis=None, dtype=None, out=None, keepdims=False):
+    """Perform a reduction using CUB.
+
+    If the specified reduction is not possible, None is returned.
+    """
+    # if import at the top level, a segfault would happen when import cupy!
+    from cupy.core._reduction import _get_axis
+    cdef bint enforce_numpy_API = False
+
+    if op > CUPY_CUB_MAX:
+        # For argmin and argmax, NumPy does not allow a tuple for axis.
+        # Also, the keepdims and dtype kwargs are not provided.
+        #
+        # For now we don't enforce these for consistency with existing CuPy
+        # non-CUB reduction behavior.
+        # https://github.com/cupy/cupy/issues/2595
+        enforce_numpy_API = False
+        if enforce_numpy_API:
+            # numpy's argmin and argmax do not support a tuple of axes
+            if sequence.PySequence_Check(axis):
+                raise TypeError(
+                    "'tuple' object cannot be interpreted as an integer")
+            if keepdims:
+                raise TypeError(
+                    "'keepdims' is an invalid keyword argument for "
+                    "argmin or argmax.")
+            if dtype is not None:
+                raise TypeError(
+                    "'dtype' is an invalid keyword argument for "
+                    "argmin or argmax.")
+        else:
+            if dtype is not None:
+                # fallback to existing non-CUB behavior
+                return None
+
+    reduce_axis, out_axis = _get_axis(axis, arr.ndim)
+    if can_use_device_reduce(op, arr.dtype, out_axis, dtype):
+        return device_reduce(arr, op, out_axis, out, keepdims)
+
+    if op > CUPY_CUB_MAX:
+        # segmented reduction not currently implemented for argmax, argmin
+        return None
+
+    if arr.flags.c_contiguous:
+        order = 'C'
+    elif arr.flags.f_contiguous:
+        order = 'F'
+    else:
+        order = None
+
+    if can_use_device_reduce(op, arr.dtype, out_axis, dtype):
+        return device_reduce(arr, op, out_axis, out, keepdims)
+    return None
+
+
 def _get_dtype_id(dtype):
     if dtype == numpy.int8:
         ret = CUPY_CUB_INT8
diff --git a/tests/cupy_tests/core_tests/test_ndarray_reduction.py b/tests/cupy_tests/core_tests/test_ndarray_reduction.py
index 71dc80b6f21..c3e44153ecb 100644
--- a/tests/cupy_tests/core_tests/test_ndarray_reduction.py
+++ b/tests/cupy_tests/core_tests/test_ndarray_reduction.py
@@ -12,6 +12,12 @@ def test_max_all(self, xp, dtype):
         a = testing.shaped_random((2, 3), xp, dtype)
         return a.max()
 
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose()
+    def test_max_all_keepdims(self, xp, dtype):
+        a = testing.shaped_random((2, 3), xp, dtype)
+        return a.max(keepdims=True)
+
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose()
     def test_max_axis_large(self, xp, dtype):
@@ -36,6 +42,18 @@ def test_max_axis2(self, xp, dtype):
         a = testing.shaped_random((2, 3, 4), xp, dtype)
         return a.max(axis=2)
 
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose()
+    def test_max_multiple_axes(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        return a.max(axis=(1, 2))
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose()
+    def test_max_multiple_axes_keepdims(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        return a.max(axis=(1, 2), keepdims=True)
+
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose()
     def test_max_nan(self, xp, dtype):
@@ -60,6 +78,12 @@ def test_min_all(self, xp, dtype):
         a = testing.shaped_random((2, 3), xp, dtype)
         return a.min()
 
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose()
+    def test_min_all_keepdims(self, xp, dtype):
+        a = testing.shaped_random((2, 3), xp, dtype)
+        return a.min(keepdims=True)
+
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose()
     def test_min_axis_large(self, xp, dtype):
@@ -84,6 +108,18 @@ def test_min_axis2(self, xp, dtype):
         a = testing.shaped_random((2, 3, 4), xp, dtype)
         return a.min(axis=2)
 
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose()
+    def test_min_multiple_axes(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        return a.min(axis=(1, 2))
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose()
+    def test_min_multiple_axes_keepdims(self, xp, dtype):
+        a = testing.shaped_random((2, 3, 4), xp, dtype)
+        return a.min(axis=(1, 2), keepdims=True)
+
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose()
     def test_min_nan(self, xp, dtype):
diff --git a/tests/cupy_tests/math_tests/test_sumprod.py b/tests/cupy_tests/math_tests/test_sumprod.py
index 6a86e6a921f..82ca6bfa53a 100644
--- a/tests/cupy_tests/math_tests/test_sumprod.py
+++ b/tests/cupy_tests/math_tests/test_sumprod.py
@@ -21,6 +21,12 @@ def test_sum_all(self, xp, dtype):
         a = testing.shaped_arange((2, 3, 4), xp, dtype)
         return a.sum()
 
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_allclose()
+    def test_sum_all_keepdims(self, xp, dtype):
+        a = testing.shaped_arange((2, 3, 4), xp, dtype)
+        return a.sum(keepdims=True)
+
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose()
     def test_external_sum_all(self, xp, dtype):
@@ -126,10 +132,19 @@ def test_sum_dtype(self, xp, src_dtype, dst_dtype):
         a = testing.shaped_arange((2, 3, 4), xp, src_dtype)
         return a.sum(dtype=dst_dtype)
 
+    @testing.for_all_dtypes_combination(names=['src_dtype', 'dst_dtype'])
+    @testing.numpy_cupy_allclose()
+    def test_sum_keepdims_and_dtype(self, xp, src_dtype, dst_dtype):
+        if not xp.can_cast(src_dtype, dst_dtype):
+            return xp.array([])  # skip
+        a = testing.shaped_arange((2, 3, 4), xp, src_dtype)
+        return a.sum(axis=1, dtype=dst_dtype, keepdims=True)
+
+    @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose()
-    def test_sum_keepdims(self, xp):
-        a = testing.shaped_arange((2, 3, 4), xp)
-        return a.sum(axis=1, keepdims=True)
+    def test_sum_keepdims_multiple_axes(self, xp, dtype):
+        a = testing.shaped_arange((2, 3, 4), xp, dtype)
+        return a.sum(axis=(1, 2), keepdims=True)
 
     @testing.for_all_dtypes()
     @testing.numpy_cupy_allclose()


Code review comments:
diff_hunk: @@ -126,10 +132,19 @@ def test_sum_dtype(self, xp, src_dtype, dst_dtype):
         a = testing.shaped_arange((2, 3, 4), xp, src_dtype)
         return a.sum(dtype=dst_dtype)
 
+    @testing.for_all_dtypes_combination(names=['src_dtype', 'dst_dtype'])
+    @testing.numpy_cupy_allclose()
+    def test_sum_keepdims_and_dtype(self, xp, src_dtype, dst_dtype):
+        if not xp.can_cast(src_dtype, dst_dtype):
+            return xp.array([])  # skip
+        a = testing.shaped_arange((2, 3, 4), xp, src_dtype)
+        return a.sum(axis=1, dtype=dst_dtype, keepdims=True)
path: tests/cupy_tests/math_tests/test_sumprod.py
text: @user1:
But this case won’t use CUB to do reduction.

@author:
okay, perhaps I can change it to just set `axis=2` then.