Repository: cupy/cupy

Issue:
[FEA] Implement cupy.argwhere
A CuPy equivalent of numpy.argwhere would be a nice addition to the library. This has been mentioned in #2783 , but it's likely worth breaking out into a separate issue for tracking purposes.


Discussion:
@suryadwar if you're still interested, you may be able to implement this feature by combining `cupy.tranpose` and `cupy.nonzero`. Depending on the performance, that might be enough.
This may be even easier than that. `argwhere` is basically the single array version of `nonzero`. I'd guess that the answer already is available before we transform it into a `tuple` of 1-D arrays. In fact [this line]( https://github.com/cupy/cupy/blob/b17bdb472f4cd5f7d59ffdc7e99be0fb0cee5378/cupy/core/_routines_indexing.pyx#L76 ) suggests as much.
As usual, thanks for finding a potentially cleaner approach to a problem @jakirkham ! 😄 

Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -645,6 +645,7 @@ def isscalar(element):
 
 from cupy._sorting.search import argmax  # NOQA
 from cupy._sorting.search import argmin  # NOQA
+from cupy._sorting.search import argwhere  # NOQA
 from cupy._sorting.search import flatnonzero  # NOQA
 from cupy._sorting.search import nanargmax  # NOQA
 from cupy._sorting.search import nanargmin  # NOQA
diff --git a/cupy/_sorting/search.py b/cupy/_sorting/search.py
--- a/cupy/_sorting/search.py
+++ b/cupy/_sorting/search.py
@@ -3,8 +3,11 @@
 from cupy.core import fusion
 from cupy import util
 
+from cupy.core import _routines_indexing as _indexing
 from cupy.core import _routines_statistics as _statistics
 
+import warnings
+
 
 def argmax(a, axis=None, dtype=None, out=None, keepdims=False):
     """Returns the indices of the maximum along an axis.
@@ -212,6 +215,30 @@ def where(condition, x=None, y=None):
     return _where_ufunc(condition.astype('?'), x, y)
 
 
+def argwhere(a):
+    """Return the indices of the elements that are non-zero.
+
+    Returns a (N, ndim) dimantional array containing the
+    indices of the non-zero elements. Where `N` is number of
+    non-zero elements and `ndim` is dimention of the given array.
+
+    Args:
+        a (cupy.ndarray): array
+
+    Returns:
+        cupy.ndarray: Indices of elements that are non-zero.
+
+    .. seealso:: :func:`numpy.argwhere`
+
+    """
+    util.check_array(a, arg_name='a')
+    if a.ndim == 0:
+        warnings.warn(
+            'calling argwhere on 0d arrays is deprecated',
+            DeprecationWarning)
+    return _indexing._ndarray_argwhere(a)
+
+
 # This is to allow using the same kernels for all dtypes, ints & floats
 # as nan is a special case
 _preamble = '''


Test patch:
diff --git a/tests/cupy_tests/sorting_tests/test_search.py b/tests/cupy_tests/sorting_tests/test_search.py
--- a/tests/cupy_tests/sorting_tests/test_search.py
+++ b/tests/cupy_tests/sorting_tests/test_search.py
@@ -278,6 +278,34 @@ def test_flatnonzero(self, xp, dtype):
         return xp.flatnonzero(array)
 
 
+@testing.parameterize(
+    {'array': numpy.random.randint(0, 2, (20,))},
+    {'array': numpy.random.randn(3, 2, 4)},
+    {'array': numpy.empty((0,))},
+    {'array': numpy.empty((0, 2))},
+    {'array': numpy.empty((0, 2, 0))},
+)
+@testing.gpu
+class TestArgwhere(unittest.TestCase):
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_argwhere(self, xp, dtype):
+        array = xp.array(self.array, dtype=dtype)
+        return xp.argwhere(array)
+
+
+@testing.parameterize(
+    {'array': cupy.array(1)},
+)
+@testing.gpu
+class TestArgwhereZeroDimension(unittest.TestCase):
+
+    def test_argwhere(self):
+        with testing.assert_warns(DeprecationWarning):
+            return cupy.nonzero(self.array)
+
+
 @testing.gpu
 class TestNanArgMin(unittest.TestCase):