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
```



Discussion:
If I'm understanding correctly, the preferred behavior is for CuPy to return a (1, self.shape[0]) in this case? I'm happy to put together a pull request.
No, CuPy should return a zero-dimensional array `cupy.array(value)`, completing the computation on the GPU side as described [here](https://docs.cupy.dev/en/stable/reference/difference.html#reduction-methods). It would be nice if you post a pull request to fix it.

Resolving patch:
diff --git a/cupyx/scipy/sparse/_data.py b/cupyx/scipy/sparse/_data.py
--- 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 @@ def power(self, n, dtype=None):
 
 
 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 @@ def _arg_min_or_max(self, axis, out, op, compare):
             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 @@ def _arg_min_or_max(self, axis, out, op, compare):
                 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


Test patch:
diff --git a/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py b/tests/cupyx_tests/scipy_tests/sparse_tests/test_csc.py
--- 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 @@ def test_max_explicit(self, xp, sp):
         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
--- 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 @@ def test_max_explicit(self, xp, sp):
         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({