Skip to content

Commit a134bdc

Browse files
authored
Get rid of call_origin in dpnp.nonzero (#1764)
* Get rid of call_origin in dpnp.nonzero * Applied review comments from #1760 * Removed dppy/label/coverage channel for coverage run * Rolled back dppy/label/coverage channel * Updated note comment in description
1 parent 46375c8 commit a134bdc

File tree

6 files changed

+158
-65
lines changed

6 files changed

+158
-65
lines changed

dpnp/dpnp_array.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,7 @@ def __copy__(self):
173173
Used if :func:`copy.copy` is called on an array. Returns a copy of the array.
174174
175175
Equivalent to ``a.copy(order="K")``.
176+
176177
"""
177178
return self.copy(order="K")
178179

@@ -911,6 +912,7 @@ def max(
911912
Return the maximum along an axis.
912913
913914
Refer to :obj:`dpnp.max` for full documentation.
915+
914916
"""
915917

916918
return dpnp.max(self, axis, out, keepdims, initial, where)
@@ -922,6 +924,7 @@ def mean(
922924
Returns the average of the array elements.
923925
924926
Refer to :obj:`dpnp.mean` for full documentation.
927+
925928
"""
926929

927930
return dpnp.mean(self, axis, dtype, out, keepdims, where=where)
@@ -938,6 +941,7 @@ def min(
938941
Return the minimum along a given axis.
939942
940943
Refer to :obj:`dpnp.min` for full documentation.
944+
941945
"""
942946

943947
return dpnp.min(self, axis, out, keepdims, initial, where)
@@ -957,7 +961,12 @@ def ndim(self):
957961
# 'newbyteorder',
958962

959963
def nonzero(self):
960-
"""Return the indices of the elements that are non-zero."""
964+
"""
965+
Return the indices of the elements that are non-zero.
966+
967+
Refer to :obj:`dpnp.nonzero` for full documentation.
968+
969+
"""
961970

962971
return dpnp.nonzero(self)
963972

@@ -1015,6 +1024,7 @@ def put(self, indices, vals, /, *, axis=None, mode="wrap"):
10151024
Puts values of an array into another array along a given axis.
10161025
10171026
For full documentation refer to :obj:`numpy.put`.
1027+
10181028
"""
10191029

10201030
return dpnp.put(self, indices, vals, axis=axis, mode=mode)
@@ -1226,6 +1236,7 @@ def std(
12261236
Returns the standard deviation of the array elements, along given axis.
12271237
12281238
Refer to :obj:`dpnp.std` for full documentation.
1239+
12291240
"""
12301241

12311242
return dpnp.std(self, axis, dtype, out, ddof, keepdims, where=where)
@@ -1261,6 +1272,7 @@ def sum(
12611272
Returns the sum along a given axis.
12621273
12631274
For full documentation refer to :obj:`dpnp.sum`.
1275+
12641276
"""
12651277

12661278
return dpnp.sum(
@@ -1278,6 +1290,7 @@ def swapaxes(self, axis1, axis2):
12781290
Interchange two axes of an array.
12791291
12801292
For full documentation refer to :obj:`numpy.swapaxes`.
1293+
12811294
"""
12821295

12831296
return dpnp.swapaxes(self, axis1=axis1, axis2=axis2)
@@ -1371,6 +1384,7 @@ def var(
13711384
Returns the variance of the array elements, along given axis.
13721385
13731386
Refer to :obj:`dpnp.var` for full documentation.
1387+
13741388
"""
13751389
return dpnp.var(self, axis, dtype, out, ddof, keepdims, where=where)
13761390

dpnp/dpnp_iface_indexing.py

Lines changed: 57 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -493,62 +493,92 @@ def indices(
493493
return res
494494

495495

496-
def nonzero(x, /):
496+
def nonzero(a):
497497
"""
498498
Return the indices of the elements that are non-zero.
499499
500+
Returns a tuple of arrays, one for each dimension of `a`,
501+
containing the indices of the non-zero elements in that
502+
dimension. The values in `a` are always tested and returned in
503+
row-major, C-style order.
504+
500505
For full documentation refer to :obj:`numpy.nonzero`.
501506
507+
Parameters
508+
----------
509+
a : {dpnp.ndarray, usm_ndarray}
510+
Input array.
511+
502512
Returns
503513
-------
504514
out : tuple[dpnp.ndarray]
505515
Indices of elements that are non-zero.
506516
507-
Limitations
508-
-----------
509-
Parameters `x` is supported as either :class:`dpnp.ndarray`
510-
or :class:`dpctl.tensor.usm_ndarray`.
511-
Otherwise the function will be executed sequentially on CPU.
512-
Input array data types are limited by supported DPNP :ref:`Data types`.
513-
514517
See Also
515518
--------
516519
:obj:`dpnp.flatnonzero` : Return indices that are non-zero in
517520
the flattened version of the input array.
521+
:obj:`dpnp.ndarray.nonzero` : Equivalent ndarray method.
518522
:obj:`dpnp.count_nonzero` : Counts the number of non-zero elements
519523
in the input array.
520524
521525
Notes
522526
-----
523527
While the nonzero values can be obtained with ``a[nonzero(a)]``, it is
524-
recommended to use ``x[x.astype(bool)]`` or ``x[x != 0]`` instead, which
528+
recommended to use ``a[a.astype(bool)]`` or ``a[a != 0]`` instead, which
525529
will correctly handle 0-d arrays.
526530
527531
Examples
528532
--------
529533
>>> import dpnp as np
530534
>>> x = np.array([[3, 0, 0], [0, 4, 0], [5, 6, 0]])
531-
>>> out = np.nonzero(x)
532-
>>> for arr in out:
533-
>>> [i for i in arr]
534-
[0, 1, 2, 2]
535-
[0, 1, 0, 1]
536-
537-
>>> x2 = np.array([3, 0, 0, 0, 4, 0, 5, 6, 0])
538-
>>> out2 = np.nonzero(x2)
539-
>>> for arr in out2:
540-
>>> [i for i in arr]
541-
[0, 4, 6, 7]
535+
>>> x
536+
array([[3, 0, 0],
537+
[0, 4, 0],
538+
[5, 6, 0]])
539+
>>> np.nonzero(x)
540+
(array([0, 1, 2, 2]), array([0, 1, 0, 1]))
541+
542+
>>> x[np.nonzero(x)]
543+
array([3, 4, 5, 6])
544+
>>> np.stack(np.nonzero(x)).T
545+
array([[0, 0],
546+
[1, 1],
547+
[2, 0],
548+
[2, 1]])
549+
550+
A common use for ``nonzero`` is to find the indices of an array, where
551+
a condition is ``True.`` Given an array `a`, the condition `a` > 3 is
552+
a boolean array and since ``False`` is interpreted as ``0``,
553+
``np.nonzero(a > 3)`` yields the indices of the `a` where the condition is
554+
true.
555+
556+
>>> a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
557+
>>> a > 3
558+
array([[False, False, False],
559+
[ True, True, True],
560+
[ True, True, True]])
561+
>>> np.nonzero(a > 3)
562+
(array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))
563+
564+
Using this result to index `a` is equivalent to using the mask directly:
565+
566+
>>> a[np.nonzero(a > 3)]
567+
array([4, 5, 6, 7, 8, 9])
568+
>>> a[a > 3] # prefer this spelling
569+
array([4, 5, 6, 7, 8, 9])
570+
571+
``nonzero`` can also be called as a method of the array.
572+
573+
>>> (a > 3).nonzero()
574+
(array([1, 1, 1, 2, 2, 2]), array([0, 1, 2, 0, 1, 2]))
542575
543576
"""
544577

545-
if dpnp.is_supported_array_type(x):
546-
usx_x = dpnp.get_usm_ndarray(x)
547-
return tuple(
548-
dpnp_array._create_from_usm_ndarray(y) for y in dpt.nonzero(usx_x)
549-
)
550-
551-
return call_origin(numpy.nonzero, x)
578+
usx_a = dpnp.get_usm_ndarray(a)
579+
return tuple(
580+
dpnp_array._create_from_usm_ndarray(y) for y in dpt.nonzero(usx_a)
581+
)
552582

553583

554584
def place(x, mask, vals, /):

dpnp/dpnp_iface_searching.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -299,16 +299,16 @@ def where(condition, x=None, y=None, /):
299299
Parameters
300300
----------
301301
condition : {dpnp.ndarray, usm_ndarray}
302-
Where True, yield `x`, otherwise yield `y`.
302+
When ``True``, yield `x`, otherwise yield `y`.
303303
x, y : {dpnp.ndarray, usm_ndarray, scalar}, optional
304304
Values from which to choose. `x`, `y` and `condition` need to be
305305
broadcastable to some shape.
306306
307307
Returns
308308
-------
309309
y : dpnp.ndarray
310-
An array with elements from `x` where `condition` is True, and elements
311-
from `y` elsewhere.
310+
An array with elements from `x` when `condition` is ``True``, and
311+
elements from `y` elsewhere.
312312
313313
See Also
314314
--------

tests/test_indexing.py

Lines changed: 63 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,69 @@ def test_indexing_array_negative_strides(self):
8585
assert_array_equal(arr, 10.0)
8686

8787

88+
class TestNonzero:
89+
@pytest.mark.parametrize("list_val", [[], [0], [1]])
90+
def test_trivial(self, list_val):
91+
np_res = numpy.nonzero(numpy.array(list_val))
92+
dpnp_res = dpnp.nonzero(dpnp.array(list_val))
93+
assert_array_equal(np_res, dpnp_res)
94+
95+
@pytest.mark.parametrize("val", [0, 1])
96+
def test_0d(self, val):
97+
assert_raises(ValueError, dpnp.nonzero, dpnp.array(val))
98+
assert_raises(ValueError, dpnp.nonzero, dpnp.array(val))
99+
100+
@pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
101+
def test_1d(self, dtype):
102+
a = numpy.array([1, 0, 2, -1, 0, 0, 8], dtype=dtype)
103+
ia = dpnp.array(a)
104+
105+
np_res = numpy.nonzero(a)
106+
dpnp_res = dpnp.nonzero(ia)
107+
assert_array_equal(np_res, dpnp_res)
108+
109+
@pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
110+
def test_2d(self, dtype):
111+
a = numpy.array([[0, 1, 0], [2, 0, 3]], dtype=dtype)
112+
ia = dpnp.array(a)
113+
114+
np_res = numpy.nonzero(a)
115+
dpnp_res = dpnp.nonzero(ia)
116+
assert_array_equal(np_res, dpnp_res)
117+
118+
a = numpy.eye(3, dtype=dtype)
119+
ia = dpnp.eye(3, dtype=dtype)
120+
121+
np_res = numpy.nonzero(a)
122+
dpnp_res = dpnp.nonzero(ia)
123+
assert_array_equal(np_res, dpnp_res)
124+
125+
def test_sparse(self):
126+
for i in range(20):
127+
a = numpy.zeros(200, dtype=bool)
128+
a[i::20] = True
129+
ia = dpnp.array(a)
130+
131+
np_res = numpy.nonzero(a)
132+
dpnp_res = dpnp.nonzero(ia)
133+
assert_array_equal(np_res, dpnp_res)
134+
135+
a = numpy.zeros(400, dtype=bool)
136+
a[10 + i : 20 + i] = True
137+
a[20 + i * 2] = True
138+
ia = dpnp.array(a)
139+
140+
np_res = numpy.nonzero(a)
141+
dpnp_res = dpnp.nonzero(ia)
142+
assert_array_equal(np_res, dpnp_res)
143+
144+
@pytest.mark.parametrize("dtype", get_all_dtypes())
145+
def test_array_method(self, dtype):
146+
a = numpy.array([[1, 0, 0], [4, 0, 6]], dtype=dtype)
147+
ia = dpnp.array(a)
148+
assert_array_equal(a.nonzero(), ia.nonzero())
149+
150+
88151
class TestPutAlongAxis:
89152
@pytest.mark.parametrize(
90153
"arr_dt", get_all_dtypes(no_bool=True, no_none=True)
@@ -371,40 +434,6 @@ def test_indices(dimension, dtype, sparse):
371434
assert_array_equal(Xnp, X)
372435

373436

374-
@pytest.mark.parametrize(
375-
"array",
376-
[
377-
[],
378-
[[0, 0], [0, 0]],
379-
[[1, 0], [1, 0]],
380-
[[1, 2], [3, 4]],
381-
[[0, 1, 2], [3, 0, 5], [6, 7, 0]],
382-
[[0, 1, 0, 3, 0], [5, 0, 7, 0, 9]],
383-
[[[1, 2], [0, 4]], [[0, 2], [0, 1]], [[0, 0], [3, 1]]],
384-
[
385-
[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]],
386-
[[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]],
387-
],
388-
],
389-
ids=[
390-
"[]",
391-
"[[0, 0], [0, 0]]",
392-
"[[1, 0], [1, 0]]",
393-
"[[1, 2], [3, 4]]",
394-
"[[0, 1, 2], [3, 0, 5], [6, 7, 0]]",
395-
"[[0, 1, 0, 3, 0], [5, 0, 7, 0, 9]]",
396-
"[[[1, 2], [0, 4]], [[0, 2], [0, 1]], [[0, 0], [3, 1]]]",
397-
"[[[[1, 2, 3], [3, 4, 5]], [[1, 2, 3], [2, 1, 0]]], [[[1, 3, 5], [3, 1, 0]], [[0, 1, 2], [1, 3, 4]]]]",
398-
],
399-
)
400-
def test_nonzero(array):
401-
a = numpy.array(array)
402-
ia = dpnp.array(array)
403-
expected = numpy.nonzero(a)
404-
result = dpnp.nonzero(ia)
405-
assert_array_equal(expected, result)
406-
407-
408437
@pytest.mark.parametrize(
409438
"vals", [[100, 200], (100, 200)], ids=["[100, 200]", "(100, 200)"]
410439
)

tests/test_sycl_queue.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1758,6 +1758,18 @@ def test_indices(device, sparse):
17581758
assert_sycl_queue_equal(dpnp_array.sycl_queue, sycl_queue)
17591759

17601760

1761+
@pytest.mark.parametrize(
1762+
"device",
1763+
valid_devices,
1764+
ids=[device.filter_string for device in valid_devices],
1765+
)
1766+
def test_nonzero(device):
1767+
a = dpnp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device)
1768+
x = dpnp.nonzero(a)
1769+
for x_el in x:
1770+
assert_sycl_queue_equal(x_el.sycl_queue, a.sycl_queue)
1771+
1772+
17611773
@pytest.mark.parametrize(
17621774
"device",
17631775
valid_devices,

tests/test_usm_type.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -752,6 +752,14 @@ def test_indices_sparse(usm_type, sparse):
752752
assert i.usm_type == usm_type
753753

754754

755+
@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
756+
def test_nonzero(usm_type):
757+
a = dp.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], usm_type=usm_type)
758+
x = dp.nonzero(a)
759+
for x_el in x:
760+
assert x_el.usm_type == usm_type
761+
762+
755763
@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
756764
def test_clip(usm_type):
757765
x = dp.arange(10, usm_type=usm_type)

0 commit comments

Comments
 (0)