Skip to content

Commit 84c9533

Browse files
committed
Apply comments per PR review by @antonwolfy
Also fix a typo when `condition` is not an array
1 parent 447dc88 commit 84c9533

File tree

1 file changed

+38
-4
lines changed

1 file changed

+38
-4
lines changed

dpnp/dpnp_iface_indexing.py

Lines changed: 38 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,8 @@ def _take_1d_index(x, inds, axis, q, usm_type, out=None):
226226

227227
def compress(condition, a, axis=None, out=None):
228228
"""
229-
Return selected slices of an array along given axis.
229+
A copy of `a` without the slices along `axis` for which `condition` is
230+
``False``.
230231
231232
For full documentation refer to :obj:`numpy.choose`.
232233
@@ -238,9 +239,10 @@ def compress(condition, a, axis=None, out=None):
238239
the output is truncated to the length of `condition`.
239240
a : {dpnp.ndarray, usm_ndarray}
240241
Array to extract from.
241-
axis : {int}, optional
242-
Axis along which to extract slices. If `None`, works over the
242+
axis : {None, int}, optional
243+
Axis along which to extract slices. If ``None``, works over the
243244
flattened array.
245+
Default: ``None``.
244246
out : {None, dpnp.ndarray, usm_ndarray}, optional
245247
If provided, the result will be placed in this array. It should
246248
be of the appropriate shape and dtype.
@@ -253,9 +255,41 @@ def compress(condition, a, axis=None, out=None):
253255
254256
See also
255257
--------
258+
:obj:`dpnp.take` : Take elements from an array along an axis.
259+
:obj:`dpnp.choose` : Construct an array from an index array and a set of
260+
arrays to choose from.
261+
:obj:`dpnp.diag` : Extract a diagonal or construct a diagonal array.
262+
:obj:`dpnp.diagonal` : Return specified diagonals.
263+
:obj:`dpnp.select` : Return an array drawn from elements in `choicelist`,
264+
depending on conditions.
256265
:obj:`dpnp.ndarray.compress` : Equivalent method.
257266
:obj:`dpnp.extract` : Equivalent function when working on 1-D arrays.
267+
268+
Examples
269+
--------
270+
>>> import numpy as np
271+
>>> a = np.array([[1, 2], [3, 4], [5, 6]])
272+
>>> a
273+
array([[1, 2],
274+
[3, 4],
275+
[5, 6]])
276+
>>> np.compress([0, 1], a, axis=0)
277+
array([[3, 4]])
278+
>>> np.compress([False, True, True], a, axis=0)
279+
array([[3, 4],
280+
[5, 6]])
281+
>>> np.compress([False, True], a, axis=1)
282+
array([[2],
283+
[4],
284+
[6]])
285+
286+
Working on the flattened array does not return slices along an axis but
287+
selects elements.
288+
289+
>>> np.compress([False, True], a)
290+
array([2])
258291
"""
292+
259293
dpnp.check_supported_arrays_type(a)
260294
if axis is None:
261295
if a.ndim != 1:
@@ -269,7 +303,7 @@ def compress(condition, a, axis=None, out=None):
269303
condition,
270304
dtype=dpnp.bool,
271305
usm_type=a_ary.usm_type,
272-
sycl_queue=a_ary.q,
306+
sycl_queue=a_ary.sycl_queue,
273307
)
274308
else:
275309
cond_ary = dpnp.get_usm_ndarray(condition)

0 commit comments

Comments
 (0)