add choose (#533)

Author: Rubtsowa

dpnp/dparray.pyx

@@ -851,6 +851,18 @@ cdef class dparray:
 
         return min(self, axis)
 
+    """
+    -------------------------------------------------------------------------
+    Indexing
+    -------------------------------------------------------------------------
+    """
+
+    def choose(input, choices, out=None, mode='raise'):
+        """
+        Construct an array from an index array and a set of arrays to choose from.
+        """
+        return choose(input, choices, out, mode)
+
     """
     -------------------------------------------------------------------------
     Sorting

dpnp/dpnp_algo/dpnp_algo_indexing.pyx

@@ -39,6 +39,7 @@ from dpnp.dpnp_iface_counting import count_nonzero
 
 
 __all__ += [
+    "dpnp_choose",
     "dpnp_diag_indices",
     "dpnp_diagonal",
     "dpnp_fill_diagonal",
@@ -56,6 +57,13 @@ __all__ += [
 ]
 
 
+cpdef dparray dpnp_choose(input, choices):
+    res_array = dparray(len(input), dtype=choices[0].dtype)
+    for i in range(len(input)):
+        res_array[i] = (choices[input[i]])[i]
+    return res_array
+
+
 cpdef tuple dpnp_diag_indices(n, ndim):
     cdef dparray res_item = dpnp.arange(n, dtype=dpnp.int64)
 

dpnp/dpnp_iface_indexing.py

@@ -51,6 +51,7 @@
 
 
 __all__ = [
+    "choose",
     "diag_indices",
     "diag_indices_from",
     "diagonal",
@@ -69,6 +70,59 @@
 ]
 
 
+def choose(input, choices, out=None, mode='raise'):
+    """
+    Construct an array from an index array and a set of arrays to choose from.
+
+    For full documentation refer to :obj:`numpy.choose`.
+
+    See also
+    --------
+    :obj:`take_along_axis` : Preferable if choices is an array.
+    """
+    if not use_origin_backend(input):
+        if not isinstance(input, list) and not isinstance(input, dparray):
+            pass
+        elif not isinstance(choices, list):
+            pass
+        elif out is not None:
+            pass
+        elif mode != 'raise':
+            pass
+        elif isinstance(choices, list):
+            val = True
+            for i in range(len(choices)):
+                if not isinstance(choices[i], dparray):
+                    val = False
+                    break
+            if not val:
+                pass
+            else:
+                val = True
+                len_ = len(input)
+                size_ = choices[0].size
+                for i in range(len(choices)):
+                    if choices[i].size != size_ or choices[i].size != len_:
+                        val = False
+                        break
+                if not val:
+                    pass
+                else:
+                    val = True
+                    for i in range(len_):
+                        if input[i] >= size_:
+                            val = False
+                            break
+                    if not val:
+                        pass
+                    else:
+                        return dpnp_choose(input, choices)
+        else:
+            return dpnp_choose(input, choices)
+
+    return call_origin(numpy.choose, input, choices, out, mode)
+
+
 def diag_indices(n, ndim=2):
     """
     Return the indices to access the main diagonal of an array.

tests/skipped_tests.tbl

@@ -582,14 +582,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_by_scalar
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_index_range_overflow
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_no_axis
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_broadcast
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_broadcast2
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_broadcast_fail
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_clip
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_wrap
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_raise
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_unknown_clip
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_1D_choicelist
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_choicelist_condlist_broadcast

tests/skipped_tests_gpu.tbl

@@ -735,14 +735,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_by_scalar
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_index_range_overflow
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_take_no_axis
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_broadcast
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_broadcast2
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_broadcast_fail
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_clip
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_choose_wrap
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_raise
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestChoose::test_unknown_clip
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_1D_choicelist
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_choicelist_condlist_broadcast

tests/test_indexing.py

@@ -5,6 +5,19 @@
 import numpy
 
 
+def test_choose():
+    a = numpy.r_[:4]
+    ia = dpnp.array(a)
+    b = numpy.r_[-4:0]
+    ib = dpnp.array(b)
+    c = numpy.r_[100:500:100]
+    ic = dpnp.array(c)
+
+    expected = numpy.choose([0, 0, 0, 0], [a, b, c])
+    result = dpnp.choose([0, 0, 0, 0], [ia, ib, ic])
+    numpy.testing.assert_array_equal(expected, result)
+
+
 @pytest.mark.parametrize("offset",
                          [0, 1],
                          ids=['0', '1'])