Merge pull request #2872 from niboshi/argminmax-dtype

emcastillo · Emilio Castillo · commit 360a25a96ce1 · 2019-12-26T07:10:22.000Z
Fix `argmin`/`argmax` `dtype` argument
diff --git a/cupy/core/_routines_statistics.pyx b/cupy/core/_routines_statistics.pyx
@@ -225,27 +225,27 @@ nanmax = create_reduction_func(
 
 cdef _argmin = create_reduction_func(
     'cupy_argmin',
-    ('?->q', 'B->q', 'h->q', 'H->q', 'i->q', 'I->q', 'l->q', 'L->q',
-     'q->q', 'Q->q',
-     ('e->q', (None, 'my_argmin_float(a, b)', None, None)),
-     ('f->q', (None, 'my_argmin_float(a, b)', None, None)),
-     ('d->q', (None, 'my_argmin_float(a, b)', None, None)),
-     ('F->q', (None, 'my_argmin_complex(a, b)', None, None)),
-     ('D->q', (None, 'my_argmin_complex(a, b)', None, None))),
+    tuple(['{}->{}'.format(d, r) for r in 'qlihb' for d in '?BhHiIlLqQ'])
+    + (
+        ('e->q', (None, 'my_argmin_float(a, b)', None, None)),
+        ('f->q', (None, 'my_argmin_float(a, b)', None, None)),
+        ('d->q', (None, 'my_argmin_float(a, b)', None, None)),
+        ('F->q', (None, 'my_argmin_complex(a, b)', None, None)),
+        ('D->q', (None, 'my_argmin_complex(a, b)', None, None))),
     ('min_max_st<type_in0_raw>(in0, _J)', 'my_argmin(a, b)', 'out0 = a.index',
      'min_max_st<type_in0_raw>'),
     None, _min_max_preamble)
 
 
 cdef _argmax = create_reduction_func(
     'cupy_argmax',
-    ('?->q', 'B->q', 'h->q', 'H->q', 'i->q', 'I->q', 'l->q', 'L->q',
-     'q->q', 'Q->q',
-     ('e->q', (None, 'my_argmax_float(a, b)', None, None)),
-     ('f->q', (None, 'my_argmax_float(a, b)', None, None)),
-     ('d->q', (None, 'my_argmax_float(a, b)', None, None)),
-     ('F->q', (None, 'my_argmax_complex(a, b)', None, None)),
-     ('D->q', (None, 'my_argmax_complex(a, b)', None, None))),
+    tuple(['{}->{}'.format(d, r) for r in 'qlihb' for d in '?BhHiIlLqQ'])
+    + (
+        ('e->q', (None, 'my_argmax_float(a, b)', None, None)),
+        ('f->q', (None, 'my_argmax_float(a, b)', None, None)),
+        ('d->q', (None, 'my_argmax_float(a, b)', None, None)),
+        ('F->q', (None, 'my_argmax_complex(a, b)', None, None)),
+        ('D->q', (None, 'my_argmax_complex(a, b)', None, None))),
     ('min_max_st<type_in0_raw>(in0, _J)', 'my_argmax(a, b)', 'out0 = a.index',
      'min_max_st<type_in0_raw>'),
     None, _min_max_preamble)
diff --git a/cupy/core/core.pyx b/cupy/core/core.pyx
@@ -728,6 +728,14 @@ cdef class ndarray:
                          keepdims=False):
         """Returns the indices of the maximum along a given axis.
 
+        .. note::
+           ``dtype`` and ``keepdim`` arguments are specific to CuPy. They are
+           not in NumPy.
+
+        .. note::
+           ``axis`` argument accepts a tuple of ints, but this is specific to
+           CuPy. NumPy does not support it.
+
         .. seealso::
            :func:`cupy.argmax` for full documentation,
            :meth:`numpy.ndarray.argmax`
@@ -749,6 +757,14 @@ cdef class ndarray:
                          keepdims=False):
         """Returns the indices of the minimum along a given axis.
 
+        .. note::
+           ``dtype`` and ``keepdim`` arguments are specific to CuPy. They are
+           not in NumPy.
+
+        .. note::
+           ``axis`` argument accepts a tuple of ints, but this is specific to
+           CuPy. NumPy does not support it.
+
         .. seealso::
            :func:`cupy.argmin` for full documentation,
            :meth:`numpy.ndarray.argmin`
diff --git a/cupy/sorting/search.py b/cupy/sorting/search.py
@@ -17,6 +17,14 @@ def argmax(a, axis=None, dtype=None, out=None, keepdims=False):
     Returns:
         cupy.ndarray: The indices of the maximum of ``a`` along an axis.
 
+    .. note::
+       ``dtype`` and ``keepdim`` arguments are specific to CuPy. They are
+       not in NumPy.
+
+    .. note::
+       ``axis`` argument accepts a tuple of ints, but this is specific to
+       CuPy. NumPy does not support it.
+
     .. seealso:: :func:`numpy.argmax`
 
     """
@@ -42,6 +50,14 @@ def argmin(a, axis=None, dtype=None, out=None, keepdims=False):
     Returns:
         cupy.ndarray: The indices of the minimum of ``a`` along an axis.
 
+    .. note::
+       ``dtype`` and ``keepdim`` arguments are specific to CuPy. They are
+       not in NumPy.
+
+    .. note::
+       ``axis`` argument accepts a tuple of ints, but this is specific to
+       CuPy. NumPy does not support it.
+
     .. seealso:: :func:`numpy.argmin`
 
     """
diff --git a/tests/cupy_tests/sorting_tests/test_search.py b/tests/cupy_tests/sorting_tests/test_search.py
@@ -2,6 +2,7 @@
 
 import numpy
 
+import cupy
 from cupy import testing
 
 
@@ -153,6 +154,30 @@ def test_argmin_zero_size_axis1(self, xp, dtype):
         return a.argmin(axis=1)
 
 
+@testing.gpu
+@testing.parameterize(*testing.product({
+    'func': ['argmin', 'argmax'],
+    'is_module': [True, False],
+    'shape': [(3, 4), ()],
+}))
+class TestArgMinMaxDtype(unittest.TestCase):
+
+    @testing.for_dtypes(
+        dtypes=[numpy.int8, numpy.int16, numpy.int32, numpy.int64],
+        name='result_dtype')
+    @testing.for_all_dtypes(name='in_dtype')
+    def test_argminmax_dtype(self, in_dtype, result_dtype):
+        a = testing.shaped_random(self.shape, cupy, in_dtype)
+        if self.is_module:
+            func = getattr(cupy, self.func)
+            y = func(a, dtype=result_dtype)
+        else:
+            func = getattr(a, self.func)
+            y = func(dtype=result_dtype)
+        assert y.shape == ()
+        assert y.dtype == result_dtype
+
+
 @testing.parameterize(
     {'cond_shape': (2, 3, 4), 'x_shape': (2, 3, 4), 'y_shape': (2, 3, 4)},
     {'cond_shape': (4,),      'x_shape': (2, 3, 4), 'y_shape': (2, 3, 4)},
