Add compress

cupy/__init__.py

@@ -435,6 +435,7 @@ def result_type(*arrays_and_dtypes):
 from cupy.indexing.generate import unravel_index  # NOQA
 
 from cupy.indexing.indexing import choose  # NOQA
+from cupy.indexing.indexing import compress  # NOQA
 from cupy.indexing.indexing import diagonal  # NOQA
 from cupy.indexing.indexing import take  # NOQA
 from cupy.indexing.indexing import take_along_axis  # NOQA

cupy/core/_routines_indexing.pxd

@@ -10,6 +10,7 @@ cdef _ndarray_scatter_min(ndarray self, slices, value)
 cdef ndarray _ndarray_take(ndarray self, indices, axis, out)
 cdef ndarray _ndarray_put(ndarray self, indices, values, mode)
 cdef ndarray _ndarray_choose(ndarray self, choices, out, mode)
+cdef ndarray _ndarray_compress(ndarray self, condition, axis, out)
 cdef ndarray _ndarray_diagonal(ndarray self, offset, axis1, axis2)
 
 cdef ndarray _simple_getitem(ndarray a, list slice_list)

cupy/core/_routines_indexing.pyx

@@ -152,6 +152,22 @@ cdef ndarray _ndarray_choose(ndarray self, choices, out, mode):
     return out
 
 
+cdef ndarray _ndarray_compress(ndarray self, condition, axis, out):
+    a = self
+
+    if numpy.isscalar(condition):
+        raise ValueError('condition must be a 1-d array')
+
+    if not isinstance(condition, ndarray):
+        condition = core.array(condition, dtype=int)
+        if condition.ndim != 1:
+            raise ValueError('condition must be a 1-d array')
+
+    res = _ndarray_nonzero(condition)  # synchronize
+
+    return _ndarray_take(a, res[0], axis, out)
+
+
 cdef ndarray _ndarray_diagonal(ndarray self, offset, axis1, axis2):
     return _diagonal(self, offset, axis1, axis2)
 

cupy/core/core.pxd

@@ -43,6 +43,7 @@ cdef class ndarray:
     cpdef partition(self, kth, int axis=*)
     cpdef ndarray argpartition(self, kth, axis=*)
     cpdef tuple nonzero(self)
+    cpdef ndarray compress(self, condition, axis=*, out=*)
     cpdef ndarray diagonal(self, offset=*, axis1=*, axis2=*)
     cpdef ndarray max(self, axis=*, out=*, keepdims=*)
     cpdef ndarray argmax(self, axis=*, out=*, dtype=*,

cupy/core/core.pyx

@@ -724,7 +724,19 @@ cdef class ndarray:
 
         return _indexing._ndarray_nonzero(self)
 
-    # TODO(okuta): Implement compress
+    cpdef ndarray compress(self, condition, axis=None, out=None):
+        """Returns selected slices of this array along given axis.
+
+        .. warning::
+
+            This function may synchronize the device.
+
+        .. seealso::
+           :func:`cupy.compress` for full documentation,
+           :meth:`numpy.ndarray.compress`
+
+        """
+        return _indexing._ndarray_compress(self, condition, axis, out)
 
     cpdef ndarray diagonal(self, offset=0, axis1=0, axis2=1):
         """Returns a view of the specified diagonals.

cupy/indexing/indexing.py

@@ -76,7 +76,33 @@ def choose(a, choices, out=None, mode='raise'):
     return a.choose(choices, out, mode)
 
 
-# TODO(okuta): Implement compress
+def compress(condition, a, axis=None, out=None):
+    """Returns selected slices of an array along given axis.
+
+    Args:
+        condition (1-D array of bools): Array that selects which entries to
+            return. If len(condition) is less than the size of a along the
+            given axis, then output is truncated to the length of the condition
+            array.
+        a (cupy.ndarray): Array from which to extract a part.
+        axis (int): Axis along which to take slices. If None (default), work
+            on the flattened array.
+        out (cupy.ndarray): Output array. If provided, it should be of
+            appropriate shape and dtype.
+
+    Returns:
+        cupy.ndarray: A copy of a without the slices along axis for which
+            condition is false.
+
+    .. warning::
+
+            This function may synchronize the device.
+
+
+    .. seealso:: :func:`numpy.compress`
+
+    """
+    return a.compress(condition, axis, out)
 
 
 def diagonal(a, offset=0, axis1=0, axis2=1):

tests/cupy_tests/indexing_tests/test_indexing.py

@@ -41,6 +41,37 @@ 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)
 
+    @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)
+
+    @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)
+
+    @testing.for_int_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_compress_no_bool(self, xp, dtype):
+        a = testing.shaped_arange((3, 4, 5), xp)
+        b = testing.shaped_arange((3,), xp, dtype)
+        return xp.compress(b, a, axis=1)
+
+    @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)
+
+    @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)
+
     @testing.for_all_dtypes()
     @testing.numpy_cupy_array_equal()
     def test_diagonal(self, xp, dtype):