Add CuPy compress

Would be useful to have an compress implementation for CuPy arrays.

[grlee77]

Implementation can sometimes be tricky because you often need some familiarity with NumPy/C to determine what is going on there and then some familiarity with CUDA/C++ to be able to implement an equivalent function for CuPy.

In the case of compress, you could check whether maybe it can be implemented by adapting an existing function like _take as defined in cupy/core/_routines_indexing.pyx. (NumPy’s compress ultimately calls PyArray_Compress)

An easier, pure Python route is available in cases where a function can be defined in terms of other functions already implemented in CuPy. In the case of compress, it seems like the following the following might work:

def compress(condition, a, axis=None, out=None):
    # TODO: replace with more efficient kernel-based implementation
    return cupy.take(a, cupy.where(condition)[0], axis, out)

This is not computationally optimal as we have to call both cupy.where and cupy.take instead of just a single underlying CUDA kernel, but it is better than nothing. I think a PR of that sort could potentially be accepted if good tests are added to make sure the functionality matches what is provided in NumPy.

Tests would need to be added to the existing file tests/cupy_tests/core_tests/test_ndarray_indexing.py.

[jakirkham]

Yep, it’s still open. Thanks for giving this a go Harshan 🙂