Repository: cupy/cupy
Pull request: Fix `argmin`/`argmax` `dtype` argument
Fixes #2595
Requires #2870


Problem statement:
Signatures and behaviors of `argmax` and `argmin` are incompatible with NumPy
First, for the signature: In CuPy the signature for `argmax()` and `argmin()` is 
```python
cupy.argm*(a, axis=None, dtype=None, out=None, keepdims=False)
```
But in NumPy it's 
```python
numpy.argm*(a, axis=None, out=None)
```
That is, `dtype` and `keepdims` should be removed.

Second, for the behavior: In CuPy the `axis` could be a tuple, but in Numpy it can only be an integer:
```python
>>> import cupy as cp
>>> a = cp.arange(60).reshape(3,4,5)
>>> a.argmax(axis=(0,1))
array([11, 11, 11, 11, 11], dtype=int64)
>>> 
>>> import numpy as np
>>> a = np.arange(60).reshape(3,4,5)
>>> a.argmax(axis=(0,1))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: 'tuple' object cannot be interpreted as an integer
```

I think both are easy to fix. Just add a few guards prior to calling the actual workhorses.

Hints from issue discussion:
From my viewpoint, I don't mind at all such deviations. I think it'd help more use cases. Rather than fixing them, I think documenting them clearly is enough. But it's not my call 😄 
Although the `dtype` argument is provided in argmin/argmax in CuPy, isn't it broken for anything except `dtype=None`?.

At least when I tried it (with CUB disabled), I get a `TypeError`:
`Wrong type (<class 'numpy.float32'>) of arguments for cupy_argmin`

(whereas NumPy gives a `TypeError` stating `'dtype' is an invalid keyword argument for argmin`)


`keepdims` does work in CuPy, though, which seems useful if not strictly NumPy-compatible
For this case I think NumPy should change -- naively I'd expect `argmax`/`argmin`'s behavior should be the same as that of `max`/`min`, i.e. accepting a tuple `axis` and `keepdims` and so on. I wonder why they didn't do so.
I think `dtype` has a use case, because without it you don't have control of the retuned dtype, even if you know int64 is too much.

Issue numpy/numpy#8710: "Add `keepdims` argument to argmin and argmax"
We don't have to remove any arguments?
Then, #2872 fixes the remaining issue (`dtype` argument of `argmin/argmax`).

Commit message: Fix argmin/argmax dtype argument

Patch to review:
diff --git a/cupy/core/_routines_statistics.pyx b/cupy/core/_routines_statistics.pyx
index 6ebd5b070c9..c1f08cbfdad 100644
--- a/cupy/core/_routines_statistics.pyx
+++ b/cupy/core/_routines_statistics.pyx
@@ -206,13 +206,13 @@ 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_float(a, b)', None, None)),
-     ('D->q', (None, 'my_argmin_float(a, b)', None, None))),
+    tuple(['{}->{}'.format(d, r) for r in 'bhilq' 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_float(a, b)', None, None)),
+        ('D->q', (None, 'my_argmin_float(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)
@@ -220,13 +220,13 @@ cdef _argmin = create_reduction_func(
 
 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_float(a, b)', None, None)),
-     ('D->q', (None, 'my_argmax_float(a, b)', None, None))),
+    tuple(['{}->{}'.format(d, r) for r in 'bhilq' 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_float(a, b)', None, None)),
+        ('D->q', (None, 'my_argmax_float(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/tests/cupy_tests/sorting_tests/test_search.py b/tests/cupy_tests/sorting_tests/test_search.py
index 8f09de602f3..6373a0533f7 100644
--- 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)},


Code review comments:
diff_hunk: @@ -206,27 +206,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_float(a, b)', None, None)),
-     ('D->q', (None, 'my_argmin_float(a, b)', None, None))),
+    tuple(['{}->{}'.format(d, r) for r in 'bhilq' for d in '?BhHiIlLqQ'])
path: cupy/core/_routines_statistics.pyx
text: @user1:
I am a bit concerned with the order `'bhilq'`. Which return type is chosen when `argmin(..., dtype=None, ...)` is called? I thought we'd need to put `'q'` as the leading character to match NumPy default, but I could be wrong. (I don't have access now to test this PR.)

@author:
Thank you for pointing out! It seems you are correct. Some of the tests were failing.
Fixed in eb90e95d5e3a421d31a56c2eaec4d32a0471780d.