Get rid of call_origin in dpnp.nonzero

dpnp/dpnp_array.py

@@ -173,6 +173,7 @@ def __copy__(self):
         Used if :func:`copy.copy` is called on an array. Returns a copy of the array.
 
         Equivalent to ``a.copy(order="K")``.
+
         """
         return self.copy(order="K")
 
@@ -911,6 +912,7 @@ def max(
         Return the maximum along an axis.
 
         Refer to :obj:`dpnp.max` for full documentation.
+
         """
 
         return dpnp.max(self, axis, out, keepdims, initial, where)
@@ -922,6 +924,7 @@ def mean(
         Returns the average of the array elements.
 
         Refer to :obj:`dpnp.mean` for full documentation.
+
         """
 
         return dpnp.mean(self, axis, dtype, out, keepdims, where=where)
@@ -938,6 +941,7 @@ def min(
         Return the minimum along a given axis.
 
         Refer to :obj:`dpnp.min` for full documentation.
+
         """
 
         return dpnp.min(self, axis, out, keepdims, initial, where)
@@ -957,7 +961,12 @@ def ndim(self):
     # 'newbyteorder',
 
     def nonzero(self):
-        """Return the indices of the elements that are non-zero."""
+        """
+        Return the indices of the elements that are non-zero.
+
+        Refer to :obj:`dpnp.nonzero` for full documentation.
+
+        """
 
         return dpnp.nonzero(self)
 
@@ -1015,6 +1024,7 @@ def put(self, indices, vals, /, *, axis=None, mode="wrap"):
         Puts values of an array into another array along a given axis.
 
         For full documentation refer to :obj:`numpy.put`.
+
         """
 
         return dpnp.put(self, indices, vals, axis=axis, mode=mode)
@@ -1226,6 +1236,7 @@ def std(
         Returns the standard deviation of the array elements, along given axis.
 
         Refer to :obj:`dpnp.std` for full documentation.
+
         """
 
         return dpnp.std(self, axis, dtype, out, ddof, keepdims, where=where)
@@ -1261,6 +1272,7 @@ def sum(
         Returns the sum along a given axis.
 
         For full documentation refer to :obj:`dpnp.sum`.
+
         """
 
         return dpnp.sum(
@@ -1278,6 +1290,7 @@ def swapaxes(self, axis1, axis2):
         Interchange two axes of an array.
 
         For full documentation refer to :obj:`numpy.swapaxes`.
+
         """
 
         return dpnp.swapaxes(self, axis1=axis1, axis2=axis2)
@@ -1371,6 +1384,7 @@ def var(
         Returns the variance of the array elements, along given axis.
 
         Refer to :obj:`dpnp.var` for full documentation.
+
         """
         return dpnp.var(self, axis, dtype, out, ddof, keepdims, where=where)
 

dpnp/dpnp_iface_indexing.py

@@ -493,62 +493,92 @@ def indices(
     return res
 
 
-def nonzero(x, /):
+def nonzero(a):
     """
     Return the indices of the elements that are non-zero.
 
+    Returns a tuple of arrays, one for each dimension of `a`,
+    containing the indices of the non-zero elements in that
+    dimension. The values in `a` are always tested and returned in
+    row-major, C-style order.
+
     For full documentation refer to :obj:`numpy.nonzero`.
 
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Input array.
+
     Returns
     -------
     out : tuple[dpnp.ndarray]
         Indices of elements that are non-zero.
 
-    Limitations
-    -----------
-    Parameters `x` is supported as either :class:`dpnp.ndarray`
-    or :class:`dpctl.tensor.usm_ndarray`.
-    Otherwise the function will be executed sequentially on CPU.
-    Input array data types are limited by supported DPNP :ref:`Data types`.
-
     See Also
     --------
     :obj:`dpnp.flatnonzero` : Return indices that are non-zero in
                               the flattened version of the input array.
+    :obj:`dpnp.ndarray.nonzero` : Equivalent ndarray method.
     :obj:`dpnp.count_nonzero` : Counts the number of non-zero elements
                                 in the input array.
 
     Notes
     -----
     While the nonzero values can be obtained with ``a[nonzero(a)]``, it is
-    recommended to use ``x[x.astype(bool)]`` or ``x[x != 0]`` instead, which
+    recommended to use ``a[a.astype(bool)]`` or ``a[a != 0]`` instead, which
     will correctly handle 0-d arrays.
 
     Examples
     --------
     >>> import dpnp as np
     >>> x = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]])
-    >>> out = np.nonzero(x)
-    >>> for arr in out:
-    >>>     [i for i in arr]
-    [0, 1, 2, 2]
-    [0, 1, 0, 1]
-
-    >>> x2 = np.array([3, 0, 0, 0, 4, 0, 5, 6, 0])
-    >>> out2 = np.nonzero(x2)
-    >>> for arr in out2:
-    >>>     [i for i in arr]
-    [0, 4, 6, 7]
+    >>> x
+    array([[3, 0, 0],
+           [0, 4, 0],
+           [5, 6, 0]])
+    >>> np.nonzero(x)
+    (array([0, 1, 2, 2]), array([0, 1, 0, 1]))
+
+    >>> x[np.nonzero(x)]
+    array([3, 4, 5, 6])
+    >>> np.stack(np.nonzero(x)).T
+    array([[0, 0],
+           [1, 1],
+           [2, 0],
+           [2, 1]])
+
+    A common use for ``nonzero`` is to find the indices of an array, where
+    a condition is ``True.``  Given an array `a`, the condition `a` > 3 is
+    a boolean array and since ``False`` is interpreted as ``0``,
+    ``np.nonzero(a > 3)`` yields the indices of the `a` where the condition is
+    true.
+
+    >>> a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    >>> a > 3
+    array([[False, False, False],
+           [ True,  True,  True],
+           [ True,  True,  True]])
+    >>> np.nonzero(a > 3)
+    (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))
+
+    Using this result to index `a` is equivalent to using the mask directly:
+
+    >>> a[np.nonzero(a > 3)]
+    array([4, 5, 6, 7, 8, 9])
+    >>> a[a > 3]  # prefer this spelling
+    array([4, 5, 6, 7, 8, 9])
+
+    ``nonzero`` can also be called as a method of the array.
+
+    >>> (a > 3).nonzero()
+    (array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))
 
     """
 
-    if dpnp.is_supported_array_type(x):
-        usx_x = dpnp.get_usm_ndarray(x)
-        return tuple(
-            dpnp_array._create_from_usm_ndarray(y) for y in dpt.nonzero(usx_x)
-        )
-
-    return call_origin(numpy.nonzero, x)
+    usx_a = dpnp.get_usm_ndarray(a)
+    return tuple(
+        dpnp_array._create_from_usm_ndarray(y) for y in dpt.nonzero(usx_a)
+    )
 
 
 def place(x, mask, vals, /):

dpnp/dpnp_iface_searching.py

@@ -299,16 +299,16 @@ def where(condition, x=None, y=None, /):
     Parameters
     ----------
     condition : {dpnp.ndarray, usm_ndarray}
-        Where True, yield `x`, otherwise yield `y`.
+        When ``True``, yield `x`, otherwise yield `y`.
     x, y : {dpnp.ndarray, usm_ndarray, scalar}, optional
         Values from which to choose. `x`, `y` and `condition` need to be
         broadcastable to some shape.
 
     Returns
     -------
     y : dpnp.ndarray
-        An array with elements from `x` where `condition` is True, and elements
-        from `y` elsewhere.
+        An array with elements from `x` when `condition` is ``True``, and
+        elements from `y` elsewhere.
 
     See Also
     --------

tests/test_indexing.py

@@ -85,6 +85,69 @@ def test_indexing_array_negative_strides(self):
         assert_array_equal(arr, 10.0)
 
 
+class TestNonzero:
+    @pytest.mark.parametrize("list_val", [[], [0], [1]])
+    def test_trivial(self, list_val):
+        np_res = numpy.nonzero(numpy.array(list_val))
+        dpnp_res = dpnp.nonzero(dpnp.array(list_val))
+        assert_array_equal(np_res, dpnp_res)
+
+    @pytest.mark.parametrize("val", [0, 1])
+    def test_0d(self, val):
+        assert_raises(ValueError, dpnp.nonzero, dpnp.array(val))
+        assert_raises(ValueError, dpnp.nonzero, dpnp.array(val))
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
+    def test_1d(self, dtype):
+        a = numpy.array([1, 0, 2, -1, 0, 0, 8], dtype=dtype)
+        ia = dpnp.array(a)
+
+        np_res = numpy.nonzero(a)
+        dpnp_res = dpnp.nonzero(ia)
+        assert_array_equal(np_res, dpnp_res)
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
+    def test_2d(self, dtype):
+        a = numpy.array([[0, 1, 0], [2, 0, 3]], dtype=dtype)
+        ia = dpnp.array(a)
+
+        np_res = numpy.nonzero(a)
+        dpnp_res = dpnp.nonzero(ia)
+        assert_array_equal(np_res, dpnp_res)
+
+        a = numpy.eye(3, dtype=dtype)
+        ia = dpnp.eye(3, dtype=dtype)
+
+        np_res = numpy.nonzero(a)
+        dpnp_res = dpnp.nonzero(ia)
+        assert_array_equal(np_res, dpnp_res)
+
+    def test_sparse(self):
+        for i in range(20):
+            a = numpy.zeros(200, dtype=bool)
+            a[i::20] = True
+            ia = dpnp.array(a)
+
+            np_res = numpy.nonzero(a)
+            dpnp_res = dpnp.nonzero(ia)
+            assert_array_equal(np_res, dpnp_res)
+
+            a = numpy.zeros(400, dtype=bool)
+            a[10 + i : 20 + i] = True
+            a[20 + i * 2] = True
+            ia = dpnp.array(a)
+
+            np_res = numpy.nonzero(a)
+            dpnp_res = dpnp.nonzero(ia)
+            assert_array_equal(np_res, dpnp_res)
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    def test_array_method(self, dtype):
+        a = numpy.array([[1, 0, 0], [4, 0, 6]], dtype=dtype)
+        ia = dpnp.array(a)
+        assert_array_equal(a.nonzero(), ia.nonzero())
+
+
 class TestPutAlongAxis:
     @pytest.mark.parametrize(
         "arr_dt", get_all_dtypes(no_bool=True, no_none=True)
@@ -371,40 +434,6 @@ def test_indices(dimension, dtype, sparse):
         assert_array_equal(Xnp, X)
 
 
-@pytest.mark.parametrize(
-    "array",
-    [
-        [],
-        [[0, 0], [0, 0]],
-        [[1, 0], [1, 0]],
-        [[1, 2], [3, 4]],
-        [[0, 1, 2], [3, 0, 5], [6, 7, 0]],
-        [[0, 1, 0, 3, 0], [5, 0, 7, 0, 9]],
-        [[[1, 2], [0, 4]], [[0, 2], [0, 1]], [[0, 0], [3, 1]]],
-        [
-            [[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]],
-            [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]],
-        ],
-    ],
-    ids=[
-        "[]",
-        "[[0, 0], [0, 0]]",
-        "[[1, 0], [1, 0]]",
-        "[[1, 2], [3, 4]]",
-        "[[0, 1, 2], [3, 0, 5], [6, 7, 0]]",
-        "[[0, 1, 0, 3, 0], [5, 0, 7, 0, 9]]",
-        "[[[1, 2], [0, 4]], [[0, 2], [0, 1]], [[0, 0], [3, 1]]]",
-        "[[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]]",
-    ],
-)
-def test_nonzero(array):
-    a = numpy.array(array)
-    ia = dpnp.array(array)
-    expected = numpy.nonzero(a)
-    result = dpnp.nonzero(ia)
-    assert_array_equal(expected, result)
-
-
 @pytest.mark.parametrize(
     "vals", [[100, 200], (100, 200)], ids=["[100, 200]", "(100, 200)"]
 )

tests/test_sycl_queue.py

@@ -1758,6 +1758,18 @@ def test_indices(device, sparse):
         assert_sycl_queue_equal(dpnp_array.sycl_queue, sycl_queue)
 
 
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_nonzero(device):
+    a = dpnp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device)
+    x = dpnp.nonzero(a)
+    for x_el in x:
+        assert_sycl_queue_equal(x_el.sycl_queue, a.sycl_queue)
+
+
 @pytest.mark.parametrize(
     "device",
     valid_devices,

tests/test_usm_type.py

@@ -752,6 +752,14 @@ def test_indices_sparse(usm_type, sparse):
         assert i.usm_type == usm_type
 
 
+@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
+def test_nonzero(usm_type):
+    a = dp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], usm_type=usm_type)
+    x = dp.nonzero(a)
+    for x_el in x:
+        assert x_el.usm_type == usm_type
+
+
 @pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
 def test_clip(usm_type):
     x = dp.arange(10, usm_type=usm_type)