From 96410641995976d91cce3b1c0109b78ae416c6ad Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 16 Nov 2024 02:08:35 -0800
Subject: [PATCH 01/20] Implement dpnp.compress

---
 dpnp/dpnp_iface_indexing.py | 142 ++++++++++++++++++++++++++++++++++++
 1 file changed, 142 insertions(+)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index 9881df7c4c4a..ce0376808bde 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -40,7 +40,10 @@
 import operator
 
 import dpctl.tensor as dpt
+import dpctl.tensor._tensor_impl as ti
+import dpctl.utils as dpu
 import numpy
+from dpctl.tensor._copy_utils import _nonzero_impl
 from dpctl.tensor._numpy_helper import normalize_axis_index
 
 import dpnp
@@ -55,6 +58,7 @@
 
 __all__ = [
     "choose",
+    "compress",
     "diag_indices",
     "diag_indices_from",
     "diagonal",
@@ -155,6 +159,144 @@ def choose(x1, choices, out=None, mode="raise"):
     return call_origin(numpy.choose, x1, choices, out, mode)
 
 
+def compress(condition, a, axis=None, out=None):
+    """
+    Return selected slices of an array along given axis.
+
+    For full documentation refer to :obj:`numpy.choose`.
+
+    Parameters
+    ----------
+    condition : {array_like, dpnp.ndarray, usm_ndarray}
+        Array that selects which entries to extract. If the length of
+        `condition` is less than the size of `a` along `axis`, then
+        the output is truncated to the length of `condition`.
+    a : {dpnp.ndarray, usm_ndarray}
+        Array to extract from.
+    axis : {int}, optional
+        Axis along which to extract slices. If `None`, works over the
+        flattened array.
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
+        If provided, the result will be placed in this array. It should
+        be of the appropriate shape and dtype.
+        Default: ``None``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        A copy of the slices of `a` where `condition` is True.
+
+    See also
+    --------
+    :obj:`dpnp.ndarray.compress` : Equivalent method.
+    :obj:`dpnp.extract` : Equivalent function when working on 1-D arrays.
+    """
+    dpnp.check_supported_arrays_type(a)
+    if axis is None:
+        if a.ndim != 1:
+            a = dpnp.ravel(a)
+        axis = 0
+    else:
+        axis = normalize_axis_index(operator.index(axis), a.ndim)
+
+    a_ary = dpnp.get_usm_ndarray(a)
+    if not dpnp.is_supported_array_type(condition):
+        usm_type = a_ary.usm_type
+        q = a_ary.sycl_queue
+        cond_ary = dpnp.as_usm_ndarray(
+            condition,
+            dtype=dpnp.bool,
+            usm_type=usm_type,
+            sycl_queue=q,
+        )
+        queues_ = [q]
+        usm_types_ = [usm_type]
+    else:
+        cond_ary = dpnp.get_usm_ndarray(condition)
+        queues_ = [a_ary.sycl_queue, cond_ary.sycl_queue]
+        usm_types_ = [a_ary.usm_type, cond_ary.usm_type]
+    if not cond_ary.ndim == 1:
+        raise ValueError(
+            "`condition` must be a 1-D array or un-nested " "sequence"
+        )
+
+    res_usm_type = dpu.get_coerced_usm_type(usm_types_)
+    exec_q = dpu.get_execution_queue(queues_)
+    if exec_q is None:
+        raise dpu.ExecutionPlacementError(
+            "arrays must be allocated on the same SYCL queue"
+        )
+
+    inds = _nonzero_impl(cond_ary)  # synchronizes
+
+    res_dt = a_ary.dtype
+    ind0 = inds[0]
+    a_sh = a_ary.shape
+    axis_end = axis + 1
+    if 0 in a_sh[axis:axis_end] and ind0.size != 0:
+        raise IndexError("cannot take non-empty indices from an empty axis")
+    res_sh = a_sh[:axis] + ind0.shape + a_sh[axis_end:]
+
+    orig_out = out
+    if out is not None:
+        dpnp.check_supported_arrays_type(out)
+        out = dpnp.get_usm_ndarray(out)
+
+        if not out.flags.writable:
+            raise ValueError("provided `out` array is read-only")
+
+        if out.shape != res_sh:
+            raise ValueError(
+                "The shape of input and output arrays are inconsistent. "
+                f"Expected output shape is {res_sh}, got {out.shape}"
+            )
+
+        if res_dt != out.dtype:
+            raise ValueError(
+                f"Output array of type {res_dt} is needed, " f"got {out.dtype}"
+            )
+
+        if dpu.get_execution_queue((a_ary.sycl_queue, out.sycl_queue)) is None:
+            raise dpu.ExecutionPlacementError(
+                "Input and output allocation queues are not compatible"
+            )
+
+        if ti._array_overlap(a_ary, out):
+            # Allocate a temporary buffer to avoid memory overlapping.
+            out = dpt.empty_like(out)
+    else:
+        out = dpt.empty(
+            res_sh, dtype=res_dt, usm_type=res_usm_type, sycl_queue=exec_q
+        )
+
+    if out.size == 0:
+        return out
+
+    _manager = dpu.SequentialOrderManager[exec_q]
+    dep_evs = _manager.submitted_events
+
+    h_ev, take_ev = ti._take(
+        src=a_ary,
+        ind=inds,
+        dst=out,
+        axis_start=axis,
+        mode=0,
+        sycl_queue=exec_q,
+        depends=dep_evs,
+    )
+    _manager.add_event_pair(h_ev, take_ev)
+
+    if not (orig_out is None or orig_out is out):
+        # Copy the out data from temporary buffer to original memory
+        ht_copy_ev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
+            src=out, dst=orig_out, sycl_queue=exec_q, depends=[take_ev]
+        )
+        _manager.add_event_pair(ht_copy_ev, cpy_ev)
+        out = orig_out
+
+    return dpnp.get_result_array(out)
+
+
 def diag_indices(n, ndim=2, device=None, usm_type="device", sycl_queue=None):
     """
     Return the indices to access the main diagonal of an array.

From 466edd9c325d80eadb75f3ed56a5f4d28962d1f0 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 16 Nov 2024 02:12:39 -0800
Subject: [PATCH 02/20] Add `dpnp_array.compress` method

---
 dpnp/dpnp_array.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py
index 77f01c9a6fbe..14da3557aab9 100644
--- a/dpnp/dpnp_array.py
+++ b/dpnp/dpnp_array.py
@@ -786,7 +786,14 @@ def clip(self, min=None, max=None, out=None, **kwargs):
 
         return dpnp.clip(self, min, max, out=out, **kwargs)
 
-    # 'compress',
+    def compress(self, condition, axis=None, out=None):
+        """
+        Select slices of an array along a given axis.
+
+        Refer to :obj:`dpnp.compress` for full documentation.
+        """
+
+        return dpnp.compress(condition, self, axis=axis, out=out)
 
     def conj(self):
         """

From 07304be811b1114f79e3a907e80317de4f4b40cc Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sun, 17 Nov 2024 23:45:09 -0800
Subject: [PATCH 03/20] Break up `compress` to satisfy pylint

Also disable checks for protected access, as `compress` uses dpctl.tensor private functions
---
 dpnp/dpnp_iface_indexing.py | 142 ++++++++++++++++++------------------
 1 file changed, 73 insertions(+), 69 deletions(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index ce0376808bde..2ac605e6d4ae 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -37,6 +37,8 @@
 
 """
 
+# pylint: disable=protected-access
+
 import operator
 
 import dpctl.tensor as dpt
@@ -159,6 +161,71 @@ def choose(x1, choices, out=None, mode="raise"):
     return call_origin(numpy.choose, x1, choices, out, mode)
 
 
+def _take_1d_index(x, inds, axis, q, usm_type, out=None):
+    # arg validation assumed done by caller
+    x_sh = x.shape
+    ind0 = inds[0]
+    axis_end = axis + 1
+    if 0 in x_sh[axis:axis_end] and ind0.size != 0:
+        raise IndexError("cannot take non-empty indices from an empty axis")
+    res_sh = x_sh[:axis] + ind0.shape + x_sh[axis_end:]
+
+    orig_out = out
+    if out is not None:
+        dpnp.check_supported_arrays_type(out)
+        out = dpnp.get_usm_ndarray(out)
+
+        if not out.flags.writable:
+            raise ValueError("provided `out` array is read-only")
+
+        if out.shape != res_sh:
+            raise ValueError(
+                "The shape of input and output arrays are inconsistent. "
+                f"Expected output shape is {res_sh}, got {out.shape}"
+            )
+
+        if x.dtype != out.dtype:
+            raise ValueError(
+                f"Output array of type {x.dtype} is needed, " f"got {out.dtype}"
+            )
+
+        if dpu.get_execution_queue((q, out.sycl_queue)) is None:
+            raise dpu.ExecutionPlacementError(
+                "Input and output allocation queues are not compatible"
+            )
+
+        if ti._array_overlap(x, out):
+            # Allocate a temporary buffer to avoid memory overlapping.
+            out = dpt.empty_like(out)
+    else:
+        out = dpt.empty(res_sh, dtype=x.dtype, usm_type=usm_type, sycl_queue=q)
+
+    _manager = dpu.SequentialOrderManager[q]
+    dep_evs = _manager.submitted_events
+
+    # always use wrap mode here
+    h_ev, take_ev = ti._take(
+        src=x,
+        ind=inds,
+        dst=out,
+        axis_start=axis,
+        mode=0,
+        sycl_queue=q,
+        depends=dep_evs,
+    )
+    _manager.add_event_pair(h_ev, take_ev)
+
+    if not (orig_out is None or orig_out is out):
+        # Copy the out data from temporary buffer to original memory
+        ht_copy_ev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
+            src=out, dst=orig_out, sycl_queue=q, depends=[take_ev]
+        )
+        _manager.add_event_pair(ht_copy_ev, cpy_ev)
+        out = orig_out
+
+    return out
+
+
 def compress(condition, a, axis=None, out=None):
     """
     Return selected slices of an array along given axis.
@@ -196,8 +263,7 @@ def compress(condition, a, axis=None, out=None):
         if a.ndim != 1:
             a = dpnp.ravel(a)
         axis = 0
-    else:
-        axis = normalize_axis_index(operator.index(axis), a.ndim)
+    axis = normalize_axis_index(operator.index(axis), a.ndim)
 
     a_ary = dpnp.get_usm_ndarray(a)
     if not dpnp.is_supported_array_type(condition):
@@ -217,7 +283,7 @@ def compress(condition, a, axis=None, out=None):
         usm_types_ = [a_ary.usm_type, cond_ary.usm_type]
     if not cond_ary.ndim == 1:
         raise ValueError(
-            "`condition` must be a 1-D array or un-nested " "sequence"
+            "`condition` must be a 1-D array or un-nested sequence"
         )
 
     res_usm_type = dpu.get_coerced_usm_type(usm_types_)
@@ -227,74 +293,12 @@ def compress(condition, a, axis=None, out=None):
             "arrays must be allocated on the same SYCL queue"
         )
 
-    inds = _nonzero_impl(cond_ary)  # synchronizes
-
-    res_dt = a_ary.dtype
-    ind0 = inds[0]
-    a_sh = a_ary.shape
-    axis_end = axis + 1
-    if 0 in a_sh[axis:axis_end] and ind0.size != 0:
-        raise IndexError("cannot take non-empty indices from an empty axis")
-    res_sh = a_sh[:axis] + ind0.shape + a_sh[axis_end:]
-
-    orig_out = out
-    if out is not None:
-        dpnp.check_supported_arrays_type(out)
-        out = dpnp.get_usm_ndarray(out)
-
-        if not out.flags.writable:
-            raise ValueError("provided `out` array is read-only")
-
-        if out.shape != res_sh:
-            raise ValueError(
-                "The shape of input and output arrays are inconsistent. "
-                f"Expected output shape is {res_sh}, got {out.shape}"
-            )
-
-        if res_dt != out.dtype:
-            raise ValueError(
-                f"Output array of type {res_dt} is needed, " f"got {out.dtype}"
-            )
-
-        if dpu.get_execution_queue((a_ary.sycl_queue, out.sycl_queue)) is None:
-            raise dpu.ExecutionPlacementError(
-                "Input and output allocation queues are not compatible"
-            )
-
-        if ti._array_overlap(a_ary, out):
-            # Allocate a temporary buffer to avoid memory overlapping.
-            out = dpt.empty_like(out)
-    else:
-        out = dpt.empty(
-            res_sh, dtype=res_dt, usm_type=res_usm_type, sycl_queue=exec_q
-        )
-
-    if out.size == 0:
-        return out
+    # _nonzero_impl synchronizes and returns a tuple of usm_ndarray indices
+    inds = _nonzero_impl(cond_ary)
 
-    _manager = dpu.SequentialOrderManager[exec_q]
-    dep_evs = _manager.submitted_events
-
-    h_ev, take_ev = ti._take(
-        src=a_ary,
-        ind=inds,
-        dst=out,
-        axis_start=axis,
-        mode=0,
-        sycl_queue=exec_q,
-        depends=dep_evs,
+    return dpnp.get_result_array(
+        _take_1d_index(a_ary, inds, axis, exec_q, res_usm_type, out)
     )
-    _manager.add_event_pair(h_ev, take_ev)
-
-    if not (orig_out is None or orig_out is out):
-        # Copy the out data from temporary buffer to original memory
-        ht_copy_ev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
-            src=out, dst=orig_out, sycl_queue=exec_q, depends=[take_ev]
-        )
-        _manager.add_event_pair(ht_copy_ev, cpy_ev)
-        out = orig_out
-
-    return dpnp.get_result_array(out)
 
 
 def diag_indices(n, ndim=2, device=None, usm_type="device", sycl_queue=None):

From e0aa4108dbdef4ffda89e11de918ea60869fdb75 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sun, 17 Nov 2024 23:45:38 -0800
Subject: [PATCH 04/20] Unmute third-party tests for `compress`

---
 .../tests/third_party/cupy/indexing_tests/test_indexing.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/dpnp/tests/third_party/cupy/indexing_tests/test_indexing.py b/dpnp/tests/third_party/cupy/indexing_tests/test_indexing.py
index 3f16209ccf7f..4e858cb0acda 100644
--- a/dpnp/tests/third_party/cupy/indexing_tests/test_indexing.py
+++ b/dpnp/tests/third_party/cupy/indexing_tests/test_indexing.py
@@ -61,21 +61,18 @@ def test_take_along_axis_none_axis(self, xp):
         b = testing.shaped_random((30,), xp, dtype="int64", scale=24)
         return xp.take_along_axis(a, b, axis=None)
 
-    @pytest.mark.skip("compress() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_compress(self, xp):
         a = testing.shaped_arange((3, 4, 5), xp)
         b = xp.array([True, False, True])
         return xp.compress(b, a, axis=1)
 
-    @pytest.mark.skip("compress() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_compress_no_axis(self, xp):
         a = testing.shaped_arange((3, 4, 5), xp)
         b = xp.array([True, False, True])
         return xp.compress(b, a)
 
-    @pytest.mark.skip("compress() is not implemented yet")
     @testing.for_int_dtypes()
     @testing.numpy_cupy_array_equal()
     def test_compress_no_bool(self, xp, dtype):
@@ -83,28 +80,24 @@ def test_compress_no_bool(self, xp, dtype):
         b = testing.shaped_arange((3,), xp, dtype)
         return xp.compress(b, a, axis=1)
 
-    @pytest.mark.skip("compress() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_compress_overrun_false(self, xp):
         a = testing.shaped_arange((3,), xp)
         b = xp.array([True, False, True, False, False, False])
         return xp.compress(b, a)
 
-    @pytest.mark.skip("compress() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_compress_empty_1dim(self, xp):
         a = testing.shaped_arange((3, 4, 5), xp)
         b = xp.array([])
         return xp.compress(b, a, axis=1)
 
-    @pytest.mark.skip("compress() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_compress_empty_1dim_no_axis(self, xp):
         a = testing.shaped_arange((3, 4, 5), xp)
         b = xp.array([])
         return xp.compress(b, a)
 
-    @pytest.mark.skip("compress() is not implemented yet")
     @testing.numpy_cupy_array_equal()
     def test_compress_0dim(self, xp):
         a = xp.array(3)

From e22b968b216a25bbcc2b81b935da64e3be79a434 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 18 Nov 2024 10:14:54 -0800
Subject: [PATCH 05/20] Use `get_usm_allocations` in `compress`

---
 dpnp/dpnp_iface_indexing.py | 17 +++--------------
 1 file changed, 3 insertions(+), 14 deletions(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index 2ac605e6d4ae..3ee3ba1d3059 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -267,31 +267,20 @@ def compress(condition, a, axis=None, out=None):
 
     a_ary = dpnp.get_usm_ndarray(a)
     if not dpnp.is_supported_array_type(condition):
-        usm_type = a_ary.usm_type
-        q = a_ary.sycl_queue
         cond_ary = dpnp.as_usm_ndarray(
             condition,
             dtype=dpnp.bool,
-            usm_type=usm_type,
-            sycl_queue=q,
+            usm_type=a_ary.usm_type,
+            sycl_queue=a_ary.q,
         )
-        queues_ = [q]
-        usm_types_ = [usm_type]
     else:
         cond_ary = dpnp.get_usm_ndarray(condition)
-        queues_ = [a_ary.sycl_queue, cond_ary.sycl_queue]
-        usm_types_ = [a_ary.usm_type, cond_ary.usm_type]
     if not cond_ary.ndim == 1:
         raise ValueError(
             "`condition` must be a 1-D array or un-nested sequence"
         )
 
-    res_usm_type = dpu.get_coerced_usm_type(usm_types_)
-    exec_q = dpu.get_execution_queue(queues_)
-    if exec_q is None:
-        raise dpu.ExecutionPlacementError(
-            "arrays must be allocated on the same SYCL queue"
-        )
+    res_usm_type, exec_q = get_usm_allocations([a_ary, cond_ary])
 
     # _nonzero_impl synchronizes and returns a tuple of usm_ndarray indices
     inds = _nonzero_impl(cond_ary)

From 1f5131169cead63c13ffe6fa06b2789619f9f93d Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 18 Nov 2024 10:46:07 -0800
Subject: [PATCH 06/20] Fix bug where `out` in `compress` is dpnp_array

Also removes an unnecessary check per PR review
---
 dpnp/dpnp_iface_indexing.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index 3ee3ba1d3059..ea44da8dc797 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -170,10 +170,9 @@ def _take_1d_index(x, inds, axis, q, usm_type, out=None):
         raise IndexError("cannot take non-empty indices from an empty axis")
     res_sh = x_sh[:axis] + ind0.shape + x_sh[axis_end:]
 
-    orig_out = out
+    orig_out = None
     if out is not None:
-        dpnp.check_supported_arrays_type(out)
-        out = dpnp.get_usm_ndarray(out)
+        orig_out = out = dpnp.get_usm_ndarray(out)
 
         if not out.flags.writable:
             raise ValueError("provided `out` array is read-only")
@@ -286,7 +285,7 @@ def compress(condition, a, axis=None, out=None):
     inds = _nonzero_impl(cond_ary)
 
     return dpnp.get_result_array(
-        _take_1d_index(a_ary, inds, axis, exec_q, res_usm_type, out)
+        _take_1d_index(a_ary, inds, axis, exec_q, res_usm_type, out), out=out
     )
 
 

From fc85e73a6adb3bba399459c085c46271c536324f Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 18 Nov 2024 11:35:49 -0800
Subject: [PATCH 07/20] Apply comments per PR review by @antonwolfy

Also fix a typo when `condition` is not an array
---
 dpnp/dpnp_iface_indexing.py | 42 +++++++++++++++++++++++++++++++++----
 1 file changed, 38 insertions(+), 4 deletions(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index ea44da8dc797..55ddb1c87920 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -227,7 +227,8 @@ def _take_1d_index(x, inds, axis, q, usm_type, out=None):
 
 def compress(condition, a, axis=None, out=None):
     """
-    Return selected slices of an array along given axis.
+    A copy of `a` without the slices along `axis` for which `condition` is
+    ``False``.
 
     For full documentation refer to :obj:`numpy.choose`.
 
@@ -239,9 +240,10 @@ def compress(condition, a, axis=None, out=None):
         the output is truncated to the length of `condition`.
     a : {dpnp.ndarray, usm_ndarray}
         Array to extract from.
-    axis : {int}, optional
-        Axis along which to extract slices. If `None`, works over the
+    axis : {None, int}, optional
+        Axis along which to extract slices. If ``None``, works over the
         flattened array.
+        Default: ``None``.
     out : {None, dpnp.ndarray, usm_ndarray}, optional
         If provided, the result will be placed in this array. It should
         be of the appropriate shape and dtype.
@@ -254,9 +256,41 @@ def compress(condition, a, axis=None, out=None):
 
     See also
     --------
+    :obj:`dpnp.take` :  Take elements from an array along an axis.
+    :obj:`dpnp.choose` : Construct an array from an index array and a set of
+                         arrays to choose from.
+    :obj:`dpnp.diag` : Extract a diagonal or construct a diagonal array.
+    :obj:`dpnp.diagonal` : Return specified diagonals.
+    :obj:`dpnp.select` : Return an array drawn from elements in `choicelist`,
+                         depending on conditions.
     :obj:`dpnp.ndarray.compress` : Equivalent method.
     :obj:`dpnp.extract` : Equivalent function when working on 1-D arrays.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> a = np.array([[1, 2], [3, 4], [5, 6]])
+    >>> a
+    array([[1, 2],
+           [3, 4],
+           [5, 6]])
+    >>> np.compress([0, 1], a, axis=0)
+    array([[3, 4]])
+    >>> np.compress([False, True, True], a, axis=0)
+    array([[3, 4],
+           [5, 6]])
+    >>> np.compress([False, True], a, axis=1)
+    array([[2],
+           [4],
+           [6]])
+
+    Working on the flattened array does not return slices along an axis but
+    selects elements.
+
+    >>> np.compress([False, True], a)
+    array([2])
     """
+
     dpnp.check_supported_arrays_type(a)
     if axis is None:
         if a.ndim != 1:
@@ -270,7 +304,7 @@ def compress(condition, a, axis=None, out=None):
             condition,
             dtype=dpnp.bool,
             usm_type=a_ary.usm_type,
-            sycl_queue=a_ary.q,
+            sycl_queue=a_ary.sycl_queue,
         )
     else:
         cond_ary = dpnp.get_usm_ndarray(condition)

From 2381662be1917db04720b75b939b6c8086d08396 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 18 Nov 2024 13:26:48 -0800
Subject: [PATCH 08/20] Remove branching when `condition` is an array

Also tweaks to docstring
---
 dpnp/dpnp_iface_indexing.py | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index 55ddb1c87920..f4b8a109635f 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -227,8 +227,10 @@ def _take_1d_index(x, inds, axis, q, usm_type, out=None):
 
 def compress(condition, a, axis=None, out=None):
     """
-    A copy of `a` without the slices along `axis` for which `condition` is
-    ``False``.
+    Return selected slices of an array along given axis.
+
+    A slice of `a` is returned for each index along `axis` where `condition`
+    is ``True``.
 
     For full documentation refer to :obj:`numpy.choose`.
 
@@ -299,15 +301,13 @@ def compress(condition, a, axis=None, out=None):
     axis = normalize_axis_index(operator.index(axis), a.ndim)
 
     a_ary = dpnp.get_usm_ndarray(a)
-    if not dpnp.is_supported_array_type(condition):
-        cond_ary = dpnp.as_usm_ndarray(
-            condition,
-            dtype=dpnp.bool,
-            usm_type=a_ary.usm_type,
-            sycl_queue=a_ary.sycl_queue,
-        )
-    else:
-        cond_ary = dpnp.get_usm_ndarray(condition)
+    cond_ary = dpnp.as_usm_ndarray(
+        condition,
+        dtype=dpnp.bool,
+        usm_type=a_ary.usm_type,
+        sycl_queue=a_ary.sycl_queue,
+    )
+
     if not cond_ary.ndim == 1:
         raise ValueError(
             "`condition` must be a 1-D array or un-nested sequence"

From 49bfa0e13994df6496c73e5ca5c9be8e787fc67f Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 18 Nov 2024 14:14:40 -0800
Subject: [PATCH 09/20] Add tests for `compress`

---
 dpnp/tests/test_indexing.py | 50 +++++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py
index a3b7ddf5c349..ea0631f02c00 100644
--- a/dpnp/tests/test_indexing.py
+++ b/dpnp/tests/test_indexing.py
@@ -1,9 +1,11 @@
 import functools
 
+import dpctl
 import dpctl.tensor as dpt
 import numpy
 import pytest
 from dpctl.tensor._numpy_helper import AxisError
+from dpctl.utils import ExecutionPlacementError
 from numpy.testing import (
     assert_,
     assert_array_equal,
@@ -1333,3 +1335,51 @@ def test_error(self):
             dpnp.select([x0], [x1], default=x1)
         with pytest.raises(TypeError):
             dpnp.select([x1], [x1])
+
+
+def test_compress_basic():
+    a = dpnp.arange(16).reshape(4, 4)
+    condition = dpnp.asarray([True, False, True])
+    r = dpnp.compress(condition, a, axis=0)
+    assert_array_equal(r[0], a[0])
+    assert_array_equal(r[1], a[2])
+
+
+@pytest.mark.parametrize("dtype", get_all_dtypes())
+def test_compress_condition_all_dtypes(dtype):
+    a = dpnp.arange(10, dtype="i4")
+    condition = dpnp.tile(dpnp.asarray([0, 1], dtype=dtype), 5)
+    r = dpnp.compress(condition, a)
+    assert_array_equal(r, a[1::2])
+
+
+def test_compress_invalid_out_errors():
+    q1 = dpctl.SyclQueue()
+    q2 = dpctl.SyclQueue()
+    a = dpnp.ones(10, dtype="i4", sycl_queue=q1)
+    condition = dpnp.asarray([True], sycl_queue=q1)
+    out_bad_shape = dpnp.empty_like(a)
+    with pytest.raises(ValueError):
+        dpnp.compress(condition, a, out=out_bad_shape)
+    out_bad_queue = dpnp.empty(1, dtype="i4", sycl_queue=q2)
+    with pytest.raises(ExecutionPlacementError):
+        dpnp.compress(condition, a, out=out_bad_queue)
+    out_bad_dt = dpnp.empty(1, dtype="i8", sycl_queue=q1)
+    with pytest.raises(ValueError):
+        dpnp.compress(condition, a, out=out_bad_dt)
+    out_read_only = dpnp.empty(1, dtype="i4", sycl_queue=q1)
+    out_read_only.flags.writable = False
+    with pytest.raises(ValueError):
+        dpnp.compress(condition, a, out=out_read_only)
+
+
+def test_compress_empty_axis():
+    a = dpnp.ones((10, 0, 5), dtype="i4")
+    condition = [True, False, True]
+    r = dpnp.compress(condition, a, axis=0)
+    assert r.shape == (2, 0, 5)
+    # empty take from empty axis is permitted
+    assert dpnp.compress([False], a, axis=1).shape == (10, 0, 5)
+    # non-empty take from empty axis raises IndexError
+    with pytest.raises(IndexError):
+        dpnp.compress(condition, a, axis=1)

From e53b84eafa0cb556e2513e5ec901be29ae221724 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 18 Nov 2024 15:50:01 -0800
Subject: [PATCH 10/20] Re-use `_take_index` for `dpnp.take`

Should slightly improve efficiency by escaping an additional copy where `out` is not `None` and flattening of indices
---
 dpnp/dpnp_iface_indexing.py | 42 ++++++++++++++++---------------------
 1 file changed, 18 insertions(+), 24 deletions(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index f4b8a109635f..7763bad5fad0 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -46,6 +46,7 @@
 import dpctl.utils as dpu
 import numpy
 from dpctl.tensor._copy_utils import _nonzero_impl
+from dpctl.tensor._indexing_functions import _get_indexing_mode
 from dpctl.tensor._numpy_helper import normalize_axis_index
 
 import dpnp
@@ -161,14 +162,13 @@ def choose(x1, choices, out=None, mode="raise"):
     return call_origin(numpy.choose, x1, choices, out, mode)
 
 
-def _take_1d_index(x, inds, axis, q, usm_type, out=None):
+def _take_index(x, inds, axis, q, usm_type, out=None, mode=0):
     # arg validation assumed done by caller
     x_sh = x.shape
-    ind0 = inds[0]
     axis_end = axis + 1
-    if 0 in x_sh[axis:axis_end] and ind0.size != 0:
+    if 0 in x_sh[axis:axis_end] and inds.size != 0:
         raise IndexError("cannot take non-empty indices from an empty axis")
-    res_sh = x_sh[:axis] + ind0.shape + x_sh[axis_end:]
+    res_sh = x_sh[:axis] + inds.shape + x_sh[axis_end:]
 
     orig_out = None
     if out is not None:
@@ -202,13 +202,12 @@ def _take_1d_index(x, inds, axis, q, usm_type, out=None):
     _manager = dpu.SequentialOrderManager[q]
     dep_evs = _manager.submitted_events
 
-    # always use wrap mode here
     h_ev, take_ev = ti._take(
         src=x,
-        ind=inds,
+        ind=(inds,),
         dst=out,
         axis_start=axis,
-        mode=0,
+        mode=mode,
         sycl_queue=q,
         depends=dep_evs,
     )
@@ -319,7 +318,8 @@ def compress(condition, a, axis=None, out=None):
     inds = _nonzero_impl(cond_ary)
 
     return dpnp.get_result_array(
-        _take_1d_index(a_ary, inds, axis, exec_q, res_usm_type, out), out=out
+        _take_index(a_ary, inds[0], axis, exec_q, res_usm_type, out=out),
+        out=out,
     )
 
 
@@ -1974,8 +1974,8 @@ def take(a, indices, /, *, axis=None, out=None, mode="wrap"):
 
     """
 
-    if mode not in ("wrap", "clip"):
-        raise ValueError(f"`mode` must be 'wrap' or 'clip', but got `{mode}`.")
+    # sets mode to 0 for "wrap" and 1 for "clip", raises otherwise
+    mode = _get_indexing_mode(mode)
 
     usm_a = dpnp.get_usm_ndarray(a)
     if not dpnp.is_supported_array_type(indices):
@@ -1985,34 +1985,28 @@ def take(a, indices, /, *, axis=None, out=None, mode="wrap"):
     else:
         usm_ind = dpnp.get_usm_ndarray(indices)
 
+    res_usm_type, exec_q = get_usm_allocations([usm_a, usm_ind])
+
     a_ndim = a.ndim
     if axis is None:
-        res_shape = usm_ind.shape
-
         if a_ndim > 1:
-            # dpt.take requires flattened input array
+            # flatten input array
             usm_a = dpt.reshape(usm_a, -1)
+        axis = 0
     elif a_ndim == 0:
         axis = normalize_axis_index(operator.index(axis), 1)
-        res_shape = usm_ind.shape
     else:
         axis = normalize_axis_index(operator.index(axis), a_ndim)
-        a_sh = a.shape
-        res_shape = a_sh[:axis] + usm_ind.shape + a_sh[axis + 1 :]
-
-    if usm_ind.ndim != 1:
-        # dpt.take supports only 1-D array of indices
-        usm_ind = dpt.reshape(usm_ind, -1)
 
     if not dpnp.issubdtype(usm_ind.dtype, dpnp.integer):
         # dpt.take supports only integer dtype for array of indices
         usm_ind = dpt.astype(usm_ind, dpnp.intp, copy=False, casting="safe")
 
-    usm_res = dpt.take(usm_a, usm_ind, axis=axis, mode=mode)
+    usm_res = _take_index(
+        usm_a, usm_ind, axis, exec_q, res_usm_type, out=out, mode=mode
+    )
 
-    # need to reshape the result if shape of indices array was changed
-    result = dpnp.reshape(usm_res, res_shape)
-    return dpnp.get_result_array(result, out)
+    return dpnp.get_result_array(usm_res, out=out)
 
 
 def take_along_axis(a, indices, axis, mode="wrap"):

From 5818d101f26dd61454a46254ed276bed3cbf8de3 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 23 Nov 2024 13:56:36 -0800
Subject: [PATCH 11/20] Change error for incorrect out array dtype to
 `TypeError`

---
 dpnp/dpnp_iface_indexing.py | 20 +++++---------------
 dpnp/tests/test_indexing.py |  2 +-
 2 files changed, 6 insertions(+), 16 deletions(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index 7763bad5fad0..09360df6b5d4 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -170,9 +170,8 @@ def _take_index(x, inds, axis, q, usm_type, out=None, mode=0):
         raise IndexError("cannot take non-empty indices from an empty axis")
     res_sh = x_sh[:axis] + inds.shape + x_sh[axis_end:]
 
-    orig_out = None
     if out is not None:
-        orig_out = out = dpnp.get_usm_ndarray(out)
+        out = dpnp.get_usm_ndarray(out)
 
         if not out.flags.writable:
             raise ValueError("provided `out` array is read-only")
@@ -184,7 +183,7 @@ def _take_index(x, inds, axis, q, usm_type, out=None, mode=0):
             )
 
         if x.dtype != out.dtype:
-            raise ValueError(
+            raise TypeError(
                 f"Output array of type {x.dtype} is needed, " f"got {out.dtype}"
             )
 
@@ -213,14 +212,6 @@ def _take_index(x, inds, axis, q, usm_type, out=None, mode=0):
     )
     _manager.add_event_pair(h_ev, take_ev)
 
-    if not (orig_out is None or orig_out is out):
-        # Copy the out data from temporary buffer to original memory
-        ht_copy_ev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
-            src=out, dst=orig_out, sycl_queue=q, depends=[take_ev]
-        )
-        _manager.add_event_pair(ht_copy_ev, cpy_ev)
-        out = orig_out
-
     return out
 
 
@@ -317,10 +308,9 @@ def compress(condition, a, axis=None, out=None):
     # _nonzero_impl synchronizes and returns a tuple of usm_ndarray indices
     inds = _nonzero_impl(cond_ary)
 
-    return dpnp.get_result_array(
-        _take_index(a_ary, inds[0], axis, exec_q, res_usm_type, out=out),
-        out=out,
-    )
+    res = _take_index(a_ary, inds[0], axis, exec_q, res_usm_type, out=out)
+
+    return dpnp.get_result_array(res, out=out)
 
 
 def diag_indices(n, ndim=2, device=None, usm_type="device", sycl_queue=None):
diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py
index ea0631f02c00..96b2eac9194d 100644
--- a/dpnp/tests/test_indexing.py
+++ b/dpnp/tests/test_indexing.py
@@ -1365,7 +1365,7 @@ def test_compress_invalid_out_errors():
     with pytest.raises(ExecutionPlacementError):
         dpnp.compress(condition, a, out=out_bad_queue)
     out_bad_dt = dpnp.empty(1, dtype="i8", sycl_queue=q1)
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         dpnp.compress(condition, a, out=out_bad_dt)
     out_read_only = dpnp.empty(1, dtype="i4", sycl_queue=q1)
     out_read_only.flags.writable = False

From 9a4430c086de57bede8fc7323429e2b188da0535 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 23 Nov 2024 14:20:44 -0800
Subject: [PATCH 12/20] Move compress tests into a TestCompress class

---
 dpnp/tests/test_indexing.py | 90 ++++++++++++++++++-------------------
 1 file changed, 44 insertions(+), 46 deletions(-)

diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py
index 96b2eac9194d..c71093a4f550 100644
--- a/dpnp/tests/test_indexing.py
+++ b/dpnp/tests/test_indexing.py
@@ -1337,49 +1337,47 @@ def test_error(self):
             dpnp.select([x1], [x1])
 
 
-def test_compress_basic():
-    a = dpnp.arange(16).reshape(4, 4)
-    condition = dpnp.asarray([True, False, True])
-    r = dpnp.compress(condition, a, axis=0)
-    assert_array_equal(r[0], a[0])
-    assert_array_equal(r[1], a[2])
-
-
-@pytest.mark.parametrize("dtype", get_all_dtypes())
-def test_compress_condition_all_dtypes(dtype):
-    a = dpnp.arange(10, dtype="i4")
-    condition = dpnp.tile(dpnp.asarray([0, 1], dtype=dtype), 5)
-    r = dpnp.compress(condition, a)
-    assert_array_equal(r, a[1::2])
-
-
-def test_compress_invalid_out_errors():
-    q1 = dpctl.SyclQueue()
-    q2 = dpctl.SyclQueue()
-    a = dpnp.ones(10, dtype="i4", sycl_queue=q1)
-    condition = dpnp.asarray([True], sycl_queue=q1)
-    out_bad_shape = dpnp.empty_like(a)
-    with pytest.raises(ValueError):
-        dpnp.compress(condition, a, out=out_bad_shape)
-    out_bad_queue = dpnp.empty(1, dtype="i4", sycl_queue=q2)
-    with pytest.raises(ExecutionPlacementError):
-        dpnp.compress(condition, a, out=out_bad_queue)
-    out_bad_dt = dpnp.empty(1, dtype="i8", sycl_queue=q1)
-    with pytest.raises(TypeError):
-        dpnp.compress(condition, a, out=out_bad_dt)
-    out_read_only = dpnp.empty(1, dtype="i4", sycl_queue=q1)
-    out_read_only.flags.writable = False
-    with pytest.raises(ValueError):
-        dpnp.compress(condition, a, out=out_read_only)
-
-
-def test_compress_empty_axis():
-    a = dpnp.ones((10, 0, 5), dtype="i4")
-    condition = [True, False, True]
-    r = dpnp.compress(condition, a, axis=0)
-    assert r.shape == (2, 0, 5)
-    # empty take from empty axis is permitted
-    assert dpnp.compress([False], a, axis=1).shape == (10, 0, 5)
-    # non-empty take from empty axis raises IndexError
-    with pytest.raises(IndexError):
-        dpnp.compress(condition, a, axis=1)
+class TestCompress:
+    def test_compress_basic(self):
+        a = dpnp.arange(16).reshape(4, 4)
+        condition = dpnp.asarray([True, False, True])
+        r = dpnp.compress(condition, a, axis=0)
+        assert_array_equal(r[0], a[0])
+        assert_array_equal(r[1], a[2])
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    def test_compress_condition_all_dtypes(self, dtype):
+        a = dpnp.arange(10, dtype="i4")
+        condition = dpnp.tile(dpnp.asarray([0, 1], dtype=dtype), 5)
+        r = dpnp.compress(condition, a)
+        assert_array_equal(r, a[1::2])
+
+    def test_compress_invalid_out_errors(self):
+        q1 = dpctl.SyclQueue()
+        q2 = dpctl.SyclQueue()
+        a = dpnp.ones(10, dtype="i4", sycl_queue=q1)
+        condition = dpnp.asarray([True], sycl_queue=q1)
+        out_bad_shape = dpnp.empty_like(a)
+        with pytest.raises(ValueError):
+            dpnp.compress(condition, a, out=out_bad_shape)
+        out_bad_queue = dpnp.empty(1, dtype="i4", sycl_queue=q2)
+        with pytest.raises(ExecutionPlacementError):
+            dpnp.compress(condition, a, out=out_bad_queue)
+        out_bad_dt = dpnp.empty(1, dtype="i8", sycl_queue=q1)
+        with pytest.raises(TypeError):
+            dpnp.compress(condition, a, out=out_bad_dt)
+        out_read_only = dpnp.empty(1, dtype="i4", sycl_queue=q1)
+        out_read_only.flags.writable = False
+        with pytest.raises(ValueError):
+            dpnp.compress(condition, a, out=out_read_only)
+
+    def test_compress_empty_axis(self):
+        a = dpnp.ones((10, 0, 5), dtype="i4")
+        condition = [True, False, True]
+        r = dpnp.compress(condition, a, axis=0)
+        assert r.shape == (2, 0, 5)
+        # empty take from empty axis is permitted
+        assert dpnp.compress([False], a, axis=1).shape == (10, 0, 5)
+        # non-empty take from empty axis raises IndexError
+        with pytest.raises(IndexError):
+            dpnp.compress(condition, a, axis=1)

From b4264a8ecbb7a5d262f9f69530563f24bd8e69ce Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 23 Nov 2024 14:27:16 -0800
Subject: [PATCH 13/20] Use NumPy in compress tests

---
 dpnp/tests/test_indexing.py | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py
index c71093a4f550..e4fe840de3e9 100644
--- a/dpnp/tests/test_indexing.py
+++ b/dpnp/tests/test_indexing.py
@@ -1339,18 +1339,24 @@ def test_error(self):
 
 class TestCompress:
     def test_compress_basic(self):
+        conditions = [True, False, True]
+        a_np = numpy.arange(16).reshape(4, 4)
         a = dpnp.arange(16).reshape(4, 4)
-        condition = dpnp.asarray([True, False, True])
-        r = dpnp.compress(condition, a, axis=0)
-        assert_array_equal(r[0], a[0])
-        assert_array_equal(r[1], a[2])
+        cond_np = numpy.array(conditions)
+        cond = dpnp.array(conditions)
+        expected = numpy.compress(cond_np, a_np, axis=0)
+        result = dpnp.compress(cond, a, axis=0)
+        assert_array_equal(expected, result)
 
     @pytest.mark.parametrize("dtype", get_all_dtypes())
     def test_compress_condition_all_dtypes(self, dtype):
+        a_np = numpy.arange(10, dtype="i4")
         a = dpnp.arange(10, dtype="i4")
-        condition = dpnp.tile(dpnp.asarray([0, 1], dtype=dtype), 5)
-        r = dpnp.compress(condition, a)
-        assert_array_equal(r, a[1::2])
+        cond_np = numpy.tile(numpy.asarray([0, 1], dtype=dtype), 5)
+        cond = dpnp.tile(dpnp.asarray([0, 1], dtype=dtype), 5)
+        expected = numpy.compress(cond_np, a_np)
+        result = dpnp.compress(cond, a)
+        assert_array_equal(expected, result)
 
     def test_compress_invalid_out_errors(self):
         q1 = dpctl.SyclQueue()

From 81d81813d3ff7c9766c666e525705f07ad712939 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 23 Nov 2024 14:28:02 -0800
Subject: [PATCH 14/20] Add `no_none=True` to
 `test_compress_condition_all_dtypes`

---
 dpnp/tests/test_indexing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py
index e4fe840de3e9..e670d31d2d52 100644
--- a/dpnp/tests/test_indexing.py
+++ b/dpnp/tests/test_indexing.py
@@ -1348,7 +1348,7 @@ def test_compress_basic(self):
         result = dpnp.compress(cond, a, axis=0)
         assert_array_equal(expected, result)
 
-    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
     def test_compress_condition_all_dtypes(self, dtype):
         a_np = numpy.arange(10, dtype="i4")
         a = dpnp.arange(10, dtype="i4")

From ec299bcdcc6ea4f38f7a3e94ca785791df4b7cfd Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 23 Nov 2024 14:43:13 -0800
Subject: [PATCH 15/20] Add USM type and SYCL queue tests for `compress`

---
 dpnp/tests/test_sycl_queue.py | 21 +++++++++++++++++++++
 dpnp/tests/test_usm_type.py   | 14 ++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py
index 737c05ee915d..94a1462089df 100644
--- a/dpnp/tests/test_sycl_queue.py
+++ b/dpnp/tests/test_sycl_queue.py
@@ -2855,3 +2855,24 @@ def test_ix(device_0, device_1):
     ixgrid = dpnp.ix_(x0, x1)
     assert_sycl_queue_equal(ixgrid[0].sycl_queue, x0.sycl_queue)
     assert_sycl_queue_equal(ixgrid[1].sycl_queue, x1.sycl_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_compress(device):
+    a_np = numpy.arange(5)
+    a = dpnp.array(a_np, device=device)
+
+    cond_np = numpy.array([0, 1, 0])
+    cond = dpnp.array(cond_np, device=device)
+
+    expected = numpy.compress(cond_np, a_np, axis=None)
+    result = dpnp.compress(cond, a, axis=None)
+    assert_allclose(expected, result)
+
+    expected_queue = a.sycl_queue
+    result_queue = result.sycl_queue
+    assert_sycl_queue_equal(result_queue, expected_queue)
diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py
index 77cfb3af8297..a11ec118dd64 100644
--- a/dpnp/tests/test_usm_type.py
+++ b/dpnp/tests/test_usm_type.py
@@ -1746,3 +1746,17 @@ def test_ix(usm_type_0, usm_type_1):
     ixgrid = dp.ix_(x0, x1)
     assert ixgrid[0].usm_type == x0.usm_type
     assert ixgrid[1].usm_type == x1.usm_type
+
+
+@pytest.mark.parametrize("usm_type_a", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize(
+    "usm_type_cond", list_of_usm_types, ids=list_of_usm_types
+)
+def test_compress(usm_type_a, usm_type_cond):
+    a = dp.arange(5, usm_type=usm_type_a)
+    cond = dp.array([False, True, True], usm_type=usm_type_cond)
+    z = dp.compress(cond, a, axis=None)
+
+    assert a.usm_type == usm_type_a
+    assert cond.usm_type == usm_type_cond
+    assert z.usm_type == du.get_coerced_usm_type([usm_type_a, usm_type_cond])

From edace9ed251fb89a3cf5d76d27cb69f5295242eb Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Sat, 23 Nov 2024 15:05:41 -0800
Subject: [PATCH 16/20] More tests for compress added

---
 dpnp/tests/test_indexing.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py
index e670d31d2d52..4732b633723f 100644
--- a/dpnp/tests/test_indexing.py
+++ b/dpnp/tests/test_indexing.py
@@ -1387,3 +1387,22 @@ def test_compress_empty_axis(self):
         # non-empty take from empty axis raises IndexError
         with pytest.raises(IndexError):
             dpnp.compress(condition, a, axis=1)
+
+    def test_compress_in_overlaps_out(self):
+        conditions = [False, True, True]
+        a_np = numpy.arange(6)
+        a = dpnp.arange(6)
+        cond_np = numpy.array(conditions)
+        cond = dpnp.array(conditions)
+        out = a[2:4]
+        expected = numpy.compress(cond_np, a_np, axis=None)
+        result = dpnp.compress(cond, a, axis=None, out=out)
+        assert_array_equal(expected, result)
+        assert result is out
+        assert (a[2:4] == out).all()
+
+    def test_compress_condition_not_1d(self):
+        a = dpnp.arange(4)
+        cond = dpnp.ones((1, 4), dtype="?")
+        with pytest.raises(ValueError):
+            dpnp.compress(cond, a, axis=None)

From cd6add44a1238d1daf5f7ab3a29a1aa62794b293 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 25 Nov 2024 12:13:18 -0800
Subject: [PATCH 17/20] Docstring change per PR review

---
 dpnp/dpnp_iface_indexing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py
index 09360df6b5d4..21ec32f24009 100644
--- a/dpnp/dpnp_iface_indexing.py
+++ b/dpnp/dpnp_iface_indexing.py
@@ -244,7 +244,7 @@ def compress(condition, a, axis=None, out=None):
     Returns
     -------
     out : dpnp.ndarray
-        A copy of the slices of `a` where `condition` is True.
+        A copy of the slices of `a` where `condition` is ``True``.
 
     See also
     --------

From 2ec457e4dda0ab28a9a2f4de0a4dc372c66815f8 Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 25 Nov 2024 12:34:56 -0800
Subject: [PATCH 18/20] Integrate test for compute follows data in compress
 into test_2in_1out

---
 dpnp/tests/test_sycl_queue.py | 22 +---------------------
 1 file changed, 1 insertion(+), 21 deletions(-)

diff --git a/dpnp/tests/test_sycl_queue.py b/dpnp/tests/test_sycl_queue.py
index 94a1462089df..359f99de0482 100644
--- a/dpnp/tests/test_sycl_queue.py
+++ b/dpnp/tests/test_sycl_queue.py
@@ -718,6 +718,7 @@ def test_reduce_hypot(device):
         ),
         pytest.param("append", [1, 2, 3], [4, 5, 6]),
         pytest.param("arctan2", [-1, +1, +1, -1], [-1, -1, +1, +1]),
+        pytest.param("compress", [0, 1, 1, 0], [0, 1, 2, 3]),
         pytest.param("copysign", [0.0, 1.0, 2.0], [-1.0, 0.0, 1.0]),
         pytest.param(
             "corrcoef",
@@ -2855,24 +2856,3 @@ def test_ix(device_0, device_1):
     ixgrid = dpnp.ix_(x0, x1)
     assert_sycl_queue_equal(ixgrid[0].sycl_queue, x0.sycl_queue)
     assert_sycl_queue_equal(ixgrid[1].sycl_queue, x1.sycl_queue)
-
-
-@pytest.mark.parametrize(
-    "device",
-    valid_devices,
-    ids=[device.filter_string for device in valid_devices],
-)
-def test_compress(device):
-    a_np = numpy.arange(5)
-    a = dpnp.array(a_np, device=device)
-
-    cond_np = numpy.array([0, 1, 0])
-    cond = dpnp.array(cond_np, device=device)
-
-    expected = numpy.compress(cond_np, a_np, axis=None)
-    result = dpnp.compress(cond, a, axis=None)
-    assert_allclose(expected, result)
-
-    expected_queue = a.sycl_queue
-    result_queue = result.sycl_queue
-    assert_sycl_queue_equal(result_queue, expected_queue)

From 108f76f6aff9fc42738434c27f4f1d462b332e6f Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Mon, 25 Nov 2024 15:37:32 -0800
Subject: [PATCH 19/20] Add test for `dpnp_array.compress` and add a test for
 strided inputs to `compress`

---
 dpnp/tests/test_indexing.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/dpnp/tests/test_indexing.py b/dpnp/tests/test_indexing.py
index 4732b633723f..2bff01cf1b59 100644
--- a/dpnp/tests/test_indexing.py
+++ b/dpnp/tests/test_indexing.py
@@ -1348,6 +1348,16 @@ def test_compress_basic(self):
         result = dpnp.compress(cond, a, axis=0)
         assert_array_equal(expected, result)
 
+    def test_compress_method_basic(self):
+        conditions = [True, True, False, True]
+        a_np = numpy.arange(3 * 4).reshape(3, 4)
+        a = dpnp.arange(3 * 4).reshape(3, 4)
+        cond_np = numpy.array(conditions)
+        cond = dpnp.array(conditions)
+        expected = a_np.compress(cond_np, axis=1)
+        result = a.compress(cond, axis=1)
+        assert_array_equal(expected, result)
+
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_none=True))
     def test_compress_condition_all_dtypes(self, dtype):
         a_np = numpy.arange(10, dtype="i4")
@@ -1406,3 +1416,20 @@ def test_compress_condition_not_1d(self):
         cond = dpnp.ones((1, 4), dtype="?")
         with pytest.raises(ValueError):
             dpnp.compress(cond, a, axis=None)
+
+    def test_compress_strided(self):
+        a = dpnp.arange(20)
+        a_np = dpnp.asnumpy(a)
+        cond = dpnp.tile(dpnp.array([True, False, False, True]), 5)
+        cond_np = dpnp.asnumpy(cond)
+        result = dpnp.compress(cond, a)
+        expected = numpy.compress(cond_np, a_np)
+        assert_array_equal(result, expected)
+        # use axis keyword
+        a = dpnp.arange(50).reshape(10, 5)
+        a_np = dpnp.asnumpy(a)
+        cond = dpnp.array(dpnp.array([True, False, False, True, False]))
+        cond_np = dpnp.asnumpy(cond)
+        result = dpnp.compress(cond, a)
+        expected = numpy.compress(cond_np, a_np)
+        assert_array_equal(result, expected)

From 4ca0e30c64e1df3de3661ae6864b357181250a6e Mon Sep 17 00:00:00 2001
From: Nikita Grigorian <nikita.grigorian@intel.com>
Date: Tue, 26 Nov 2024 09:16:31 -0800
Subject: [PATCH 20/20] Refactor `test_compress` in test_usm_type.py into
 `test_2in_1out`

---
 dpnp/tests/test_usm_type.py | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/dpnp/tests/test_usm_type.py b/dpnp/tests/test_usm_type.py
index a11ec118dd64..d14604be725c 100644
--- a/dpnp/tests/test_usm_type.py
+++ b/dpnp/tests/test_usm_type.py
@@ -686,6 +686,7 @@ def test_1in_1out(func, data, usm_type):
         ),
         pytest.param("append", [1, 2, 3], [4, 5, 6]),
         pytest.param("arctan2", [-1, +1, +1, -1], [-1, -1, +1, +1]),
+        pytest.param("compress", [False, True, True], [0, 1, 2, 3, 4]),
         pytest.param("copysign", [0.0, 1.0, 2.0], [-1.0, 0.0, 1.0]),
         pytest.param("cross", [1.0, 2.0, 3.0], [4.0, 5.0, 6.0]),
         pytest.param("digitize", [0.2, 6.4, 3.0], [0.0, 1.0, 2.5, 4.0]),
@@ -1746,17 +1747,3 @@ def test_ix(usm_type_0, usm_type_1):
     ixgrid = dp.ix_(x0, x1)
     assert ixgrid[0].usm_type == x0.usm_type
     assert ixgrid[1].usm_type == x1.usm_type
-
-
-@pytest.mark.parametrize("usm_type_a", list_of_usm_types, ids=list_of_usm_types)
-@pytest.mark.parametrize(
-    "usm_type_cond", list_of_usm_types, ids=list_of_usm_types
-)
-def test_compress(usm_type_a, usm_type_cond):
-    a = dp.arange(5, usm_type=usm_type_a)
-    cond = dp.array([False, True, True], usm_type=usm_type_cond)
-    z = dp.compress(cond, a, axis=None)
-
-    assert a.usm_type == usm_type_a
-    assert cond.usm_type == usm_type_cond
-    assert z.usm_type == du.get_coerced_usm_type([usm_type_a, usm_type_cond])