Repository: cupy/cupy
Pull request: Make sparse argmin/max return a scalar array containing the index
Fixes #3640

The aforementioned issue only happened when `argmax/min` was called without an `axis` defined. Now `argmin/max` returns a scalar array with the index.

Problem statement:
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
```

Hints from issue 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.

Commit message: Remove element-by-element comparison

Patch to review:
diff --git a/cupyx/scipy/sparse/_data.py b/cupyx/scipy/sparse/_data.py
index a6047dce706..df8149c4340 100644
--- a/cupyx/scipy/sparse/_data.py
+++ b/cupyx/scipy/sparse/_data.py
@@ -121,15 +121,18 @@ def power(self, n, dtype=None):
 
 
 def _find_missing_index(ind, n):
-    for k, a in enumerate(ind):
-        if k != a:
-            return k
+    positions = cupy.arange(ind.size)
+    diff = ind != positions
+    non_zero = cupy.argwhere(diff).ravel()
 
-    k += 1
+    if non_zero.size != 0:
+        return non_zero[0]
+
+    k = ind.size + 1
     if k < n:
-        return k
+        return cupy.asarray(k)
     else:
-        return -1
+        return cupy.asarray(-1)
 
 
 class _minmax_mixin(object):
@@ -238,7 +241,7 @@ def _arg_min_or_max(self, axis, out, op, compare):
                 if compare(m, zero):
                     return mat.row[am] * mat.shape[1] + mat.col[am]
                 else:
-                    size = cupy.prod(mat.shape)
+                    size = cupy.prod(cupy.asarray(mat.shape))
                     if size == mat.nnz:
                         return am
                     else:


Code review comments:
diff_hunk: @@ -121,15 +121,18 @@ def power(self, n, dtype=None):
 
 
 def _find_missing_index(ind, n):
-    for k, a in enumerate(ind):
-        if k != a:
-            return k
+    positions = cupy.arange(ind.size)
+    diff = ind != positions
+    non_zero = cupy.argwhere(diff).ravel()
path: cupyx/scipy/sparse/_data.py
text: @user1:
`cupy.argwhere` also causes device-to-host sync because the length of return values varies depending on the value of elements of the inputs.

How about the following implementation?

```py
return cupy.where(
    diff.any(),
    diff.argmax(),
    cupy.asarray(ind.size if ind.size < n else -1))
```
```