Select (#527)

* add select

Author: Rubtsowa

dpnp/dpnp_algo/dpnp_algo_indexing.pyx

@@ -47,6 +47,7 @@ __all__ += [
     "dpnp_place",
     "dpnp_put",
     "dpnp_putmask",
+    "dpnp_select",
     "dpnp_take",
     "dpnp_tril_indices",
     "dpnp_tril_indices_from",
@@ -232,6 +233,22 @@ cpdef dpnp_putmask(dparray arr, dparray mask, dparray values):
             arr[i] = values[i % values_size]
 
 
+cpdef dparray dpnp_select(condlist, choicelist, default):
+    size_ = condlist[0].size
+    res_array = dparray(size_, dtype=choicelist[0].dtype)
+    pass_val = {a: default for a in range(size_)}
+    for i in range(len(condlist)):
+        for j in range(size_):
+            if (condlist[i])[j]:
+                res_array[j] = (choicelist[i])[j]
+                pass_val.pop(j)
+
+    for ind, val in pass_val.items():
+        res_array[ind] = val
+
+    return res_array.reshape(condlist[0].shape)
+
+
 cpdef dparray dpnp_take(dparray input, dparray indices):
     indices_size = indices.size
     res_array = dparray(indices_size, dtype=input.dtype)

dpnp/dpnp_iface_indexing.py

@@ -60,6 +60,7 @@
     "place",
     "put",
     "putmask",
+    "select",
     "take",
     "tril_indices",
     "tril_indices_from",
@@ -371,6 +372,43 @@ def putmask(arr, mask, values):
     return call_origin(numpy.putmask, arr, mask, values)
 
 
+def select(condlist, choicelist, default=0):
+    """
+    Return an array drawn from elements in choicelist, depending on conditions.
+    For full documentation refer to :obj:`numpy.select`.
+
+    Limitations
+    -----------
+    Arrays of input lists are supported as :obj:`dpnp.ndarray`.
+    Parameter ``default`` are supported only with default values.
+    """
+    if not use_origin_backend():
+        if not isinstance(condlist, list):
+            pass
+        elif not isinstance(condlist[0], dparray):
+            pass
+        elif not isinstance(choicelist, list):
+            pass
+        elif not isinstance(choicelist[0], dparray):
+            pass
+        elif len(condlist) != len(choicelist):
+            pass
+        elif len(condlist) == len(choicelist):
+            val = True
+            size_ = condlist[0].size
+            for i in range(len(condlist)):
+                if condlist[i].size != size_ or choicelist[i].size != size_:
+                    val = False
+            if not val:
+                pass
+            else:
+                return dpnp_select(condlist, choicelist, default)
+        else:
+            return dpnp_select(condlist, choicelist, default)
+
+    return call_origin(numpy.select, condlist, choicelist, default)
+
+
 def take(input, indices, axis=None, out=None, mode='raise'):
     """
     Take elements from an array.

tests/skipped_tests.tbl

@@ -599,7 +599,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_length_error
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_broadcastable
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_non_broadcastable
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_type_error_choicelist
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_type_error_condlist
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmask::test_putmask_non_equal_shape_raises
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPlace_param_0_{n_vals=0, shape=(7,)}::test_place

tests/skipped_tests_gpu.tbl

@@ -752,7 +752,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_length_error
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_broadcastable
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_non_broadcastable
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_type_error_choicelist
 tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_type_error_condlist
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmask::test_putmask_non_equal_shape_raises
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPlaceRaises_param_0_{shape=(7,)}::test_place_empty_value_error

tests/test_indexing.py

@@ -284,6 +284,24 @@ def test_putmask3(arr, mask, vals):
     numpy.testing.assert_array_equal(a, ia)
 
 
+def test_select():
+    cond_val1 = numpy.array([True, True, True, False, False, False, False, False, False, False])
+    cond_val2 = numpy.array([False, False, False, False, False, True, True, True, True, True])
+    icond_val1 = dpnp.array(cond_val1)
+    icond_val2 = dpnp.array(cond_val2)
+    condlist = [cond_val1, cond_val2]
+    icondlist = [icond_val1, icond_val2]
+    choice_val1 = numpy.full(10, -2)
+    choice_val2 = numpy.full(10, -1)
+    ichoice_val1 = dpnp.array(choice_val1)
+    ichoice_val2 = dpnp.array(choice_val2)
+    choicelist = [choice_val1, choice_val2]
+    ichoicelist = [ichoice_val1, ichoice_val2]
+    expected = numpy.select(condlist, choicelist)
+    result = dpnp.select(icondlist, ichoicelist)
+    numpy.testing.assert_array_equal(expected, result)
+
+
 @pytest.mark.parametrize("indices",
                          [[[0, 0], [0, 0]],
                           [[1, 2], [1, 2]],