Repository: cupy/cupy
Issue: Sparse argmin/argmax return scalar values on CPU

CuPy returns zero-dimensional ndarrays for reduction functions where NumPy/SciPy return scalar values to keep off synchronizing between GPU and CPU. However, CuPy's sparse matrix argmin/argmax currently return scalar values as SciPy's counterparts, causing device to host sync.

```Python
>>> import cupy
>>> import cupyx
>>> a = cupy.array([[0,1,0],[2,0,3]], dtype='f')
>>> cupyx.scipy.sparse.csc_matrix(a).argmin()
0
```


Changed files:
tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py, tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py, cupyx/scipy/sparse/_data.py

diff --git a/cupyx/scipy/sparse/_data.py b/cupyx/scipy/sparse/_data.py
index a6047dce7..1615aea2d 100644
--- a/cupyx/scipy/sparse/_data.py
+++ b/cupyx/scipy/sparse/_data.py
@@ -1,4 +1,5 @@
 import cupy
+import numpy as np
 from cupy._core import internal
 from cupy import _util
 from cupyx.scipy.sparse import _base
@@ -121,15 +122,25 @@ class _data_matrix(_base.spmatrix):
 
 
 def _find_missing_index(ind, n):
-    for k, a in enumerate(ind):
-        if k != a:
-            return k
-
-    k += 1
-    if k < n:
-        return k
+    positions = cupy.arange(ind.size)
+    diff = ind != positions
+    return cupy.where(
+        diff.any(),
+        diff.argmax(),
+        cupy.asarray(ind.size if ind.size < n else -1))
+
+
+def _non_zero_cmp(mat, am, zero, m):
+    size = np.prod(mat.shape)
+    if size == mat.nnz:
+        return am
     else:
-        return -1
+        ind = mat.row * mat.shape[1] + mat.col
+        zero_ind = _find_missing_index(ind, size)
+        return cupy.where(
+            m == zero,
+            cupy.minimum(zero_ind, am),
+            zero_ind)
 
 
 class _minmax_mixin(object):
@@ -227,7 +238,7 @@ class _minmax_mixin(object):
             if self.nnz == 0:
                 return 0
             else:
-                zero = self.dtype.type(0)
+                zero = cupy.asarray(self.dtype.type(0))
                 mat = self.tocoo()
 
                 mat.sum_duplicates()
@@ -235,19 +246,9 @@ class _minmax_mixin(object):
                 am = op(mat.data)
                 m = mat.data[am]
 
-                if compare(m, zero):
-                    return mat.row[am] * mat.shape[1] + mat.col[am]
-                else:
-                    size = cupy.prod(mat.shape)
-                    if size == mat.nnz:
-                        return am
-                    else:
-                        ind = mat.row * mat.shape[1] + mat.col
-                        zero_ind = _find_missing_index(ind, size)
-                        if m == zero:
-                            return min(zero_ind, am)
-                        else:
-                            return zero_ind
+                return cupy.where(
+                    compare(m, zero), mat.row[am] * mat.shape[1] + mat.col[am],
+                    _non_zero_cmp(mat, am, zero, m))
 
         if axis < 0:
             axis += 2
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py
index e2dbb141f..6b2c78b64 100644
--- a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py
+++ b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py
@@ -1258,21 +1258,23 @@ class TestCscMatrixScipyCompressedMinMax:
         else:
             return data.max(axis=self.axis)
 
-    @testing.numpy_cupy_array_equal(sp_name='sp')
+    @testing.numpy_cupy_array_equal(sp_name='sp', type_check=False)
     def test_argmin(self, xp, sp):
-        # TODO(takagi) Fix axis=None
-        if self.axis is None:
-            pytest.skip()
         data = self._make_data_min(xp, sp, dense=self.dense)
-        return data.argmin(axis=self.axis)
+        # Due to a SciPy bug, the argmin output is different from the expected
+        # one
+        if self.axis is None and self.dense:
+            pytest.skip()
+        return xp.array(data.argmin(axis=self.axis))
 
-    @testing.numpy_cupy_array_equal(sp_name='sp')
+    @testing.numpy_cupy_array_equal(sp_name='sp', type_check=False)
     def test_argmax(self, xp, sp):
-        # TODO(takagi) Fix axis=None
-        if self.axis is None:
-            pytest.skip()
         data = self._make_data_max(xp, sp, dense=self.dense)
-        return data.argmax(axis=self.axis)
+        # Due to a SciPy bug, the argmin output is different from the expected
+        # one
+        if self.axis is None and self.dense:
+            pytest.skip()
+        return xp.array(data.argmax(axis=self.axis))
 
 
 @testing.parameterize(*testing.product({
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
index a8cf93311..9a75ae138 100644
--- a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
+++ b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csr.py
@@ -1451,21 +1451,23 @@ class TestCsrMatrixScipyCompressedMinMax:
         else:
             return data.max(axis=self.axis)
 
-    @testing.numpy_cupy_array_equal(sp_name='sp')
+    @testing.numpy_cupy_array_equal(sp_name='sp', type_check=False)
     def test_argmin(self, xp, sp):
-        # TODO(takagi) Fix axis=None
-        if self.axis is None:
-            pytest.skip()
         data = self._make_data_min(xp, sp, dense=self.dense)
-        return data.argmin(axis=self.axis)
+        # Due to a SciPy bug, the argmin output is different from the expected
+        # one
+        if self.axis is None and self.dense:
+            pytest.skip()
+        return xp.array(data.argmin(axis=self.axis))
 
-    @testing.numpy_cupy_array_equal(sp_name='sp')
+    @testing.numpy_cupy_array_equal(sp_name='sp', type_check=False)
     def test_argmax(self, xp, sp):
-        # TODO(takagi) Fix axis=None
-        if self.axis is None:
-            pytest.skip()
         data = self._make_data_max(xp, sp, dense=self.dense)
-        return data.argmax(axis=self.axis)
+        # Due to a SciPy bug, the argmin output is different from the expected
+        # one
+        if self.axis is None and self.dense:
+            pytest.skip()
+        return xp.array(data.argmax(axis=self.axis))
 
 
 @testing.parameterize(*testing.product({