Implement dpnp.select

dpnp/dpnp_algo/dpnp_algo_indexing.pxi

@@ -38,7 +38,6 @@ and the rest of the library
 __all__ += [
     "dpnp_choose",
     "dpnp_putmask",
-    "dpnp_select",
 ]
 
 ctypedef c_dpctl.DPCTLSyclEventRef(*fptr_dpnp_choose_t)(c_dpctl.DPCTLSyclQueueRef,
@@ -102,20 +101,3 @@ cpdef dpnp_putmask(utils.dpnp_descriptor arr, utils.dpnp_descriptor mask, utils.
     for i in range(arr.size):
         if mask_flatiter[i]:
             arr_flatiter[i] = values_flatiter[i % values_size]
-
-
-cpdef utils.dpnp_descriptor dpnp_select(list condlist, list choicelist, default):
-    cdef size_t size_ = condlist[0].size
-    cdef utils.dpnp_descriptor res_array = utils_py.create_output_descriptor_py(condlist[0].shape, choicelist[0].dtype, None)
-
-    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.get_pyobj()[j] = (choicelist[i])[j]
-                pass_val.pop(j)
-
-    for ind, val in pass_val.items():
-        res_array.get_pyobj()[ind] = val
-
-    return res_array

dpnp/dpnp_iface_indexing.py

@@ -49,12 +49,11 @@
 from .dpnp_algo import (
     dpnp_choose,
     dpnp_putmask,
-    dpnp_select,
 )
 from .dpnp_array import dpnp_array
 from .dpnp_utils import (
     call_origin,
-    use_origin_backend,
+    get_usm_allocations,
 )
 
 __all__ = [
@@ -524,7 +523,7 @@ def extract(condition, a):
     :obj:`dpnp.put` : Replaces specified elements of an array with given values.
     :obj:`dpnp.copyto` : Copies values from one array to another, broadcasting
                          as necessary.
-    :obj:`dpnp.compress` : eturn selected slices of an array along given axis.
+    :obj:`dpnp.compress` : Return selected slices of an array along given axis.
     :obj:`dpnp.place` : Change elements of an array based on conditional and
                         input values.
 
@@ -1344,31 +1343,125 @@ def select(condlist, choicelist, default=0):
 
     For full documentation refer to :obj:`numpy.select`.
 
-    Limitations
-    -----------
-    Arrays of input lists are supported as :obj:`dpnp.ndarray`.
-    Parameter `default` is supported only with default values.
+    Parameters
+    ----------
+    condlist : list of bool dpnp.ndarray or usm_ndarray
+        The list of conditions which determine from which array in `choicelist`
+        the output elements are taken. When multiple conditions are satisfied,
+        the first one encountered in `condlist` is used.
+    choicelist : list of dpnp.ndarray or usm_ndarray
+        The list of arrays from which the output elements are taken. It has
+        to be of the same length as `condlist`.
+    default : {scalar, dpnp.ndarray, usm_ndarray}, optional
+        The element inserted in `output` when all conditions evaluate to
+        ``False``. Default: ``0``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        The output at position m is the m-th element of the array in
+        `choicelist` where the m-th element of the corresponding array in
+        `condlist` is ``True``.
+
+    See Also
+    --------
+    :obj:`dpnp.where : Return elements from one of two arrays depending on
+                       condition.
+    :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.compress` : Return selected slices of an array along given axis.
+    :obj:`dpnp.diag` : Extract a diagonal or construct a diagonal array.
+    :obj:`dpnp.diagonal` : Return specified diagonals.
+
+    Examples
+    --------
+    >>> import dpnp as np
+
+    Beginning with an array of integers from 0 to 5 (inclusive),
+    elements less than ``3`` are negated, elements greater than ``3``
+    are squared, and elements not meeting either of these conditions
+    (exactly ``3``) are replaced with a `default` value of ``42``.
+
+    >>> x = np.arange(6)
+    >>> condlist = [x<3, x>3]
+    >>> choicelist = [x, x**2]
+    >>> np.select(condlist, choicelist, 42)
+    array([ 0,  1,  2, 42, 16, 25])
+
+    When multiple conditions are satisfied, the first one encountered in
+    `condlist` is used.
+
+    >>> condlist = [x<=4, x>3]
+    >>> choicelist = [x, x**2]
+    >>> np.select(condlist, choicelist, 55)
+    array([ 0,  1,  2,  3,  4, 25])
+
     """
 
-    if not use_origin_backend():
-        if not isinstance(condlist, list):
-            pass
-        elif not isinstance(choicelist, list):
-            pass
-        elif len(condlist) != len(choicelist):
-            pass
-        else:
-            val = True
-            size_ = condlist[0].size
-            for cond, choice in zip(condlist, choicelist):
-                if cond.size != size_ or choice.size != size_:
-                    val = False
-            if not val:
-                pass
-            else:
-                return dpnp_select(condlist, choicelist, default).get_pyobj()
+    if len(condlist) != len(choicelist):
+        raise ValueError(
+            "list of cases must be same length as list of conditions"
+        )
+
+    if len(condlist) == 0:
+        raise ValueError("select with an empty condition list is not possible")
+
+    dpnp.check_supported_arrays_type(*condlist)
+    dpnp.check_supported_arrays_type(*choicelist)
+    dpnp.check_supported_arrays_type(
+        default, scalar_type=True, all_scalars=True
+    )
+
+    if dpnp.isscalar(default):
+        usm_type_alloc, sycl_queue_alloc = get_usm_allocations(
+            condlist + choicelist
+        )
+        dtype = dpnp.result_type(*choicelist)
+        default = dpnp.asarray(
+            default,
+            dtype=dtype,
+            usm_type=usm_type_alloc,
+            sycl_queue=sycl_queue_alloc,
+        )
+        choicelist.append(default)
+    else:
+        choicelist.append(default)
+        usm_type_alloc, sycl_queue_alloc = get_usm_allocations(
+            condlist + choicelist
+        )
+        dtype = dpnp.result_type(*choicelist)
+
+    for i, cond in enumerate(condlist):
+        if cond.dtype.type is not dpnp.bool:
+            raise TypeError(
+                f"invalid entry {i} in condlist: should be boolean ndarray"
+            )
+
+    # Convert conditions to arrays and broadcast conditions and choices
+    # as the shape is needed for the result
+    condlist = dpnp.broadcast_arrays(*condlist)
+    choicelist = dpnp.broadcast_arrays(*choicelist)
+
+    result_shape = dpnp.broadcast_arrays(condlist[0], choicelist[0])[0].shape
+
+    result = dpnp.full(
+        result_shape,
+        choicelist[-1],
+        dtype=dtype,
+        usm_type=usm_type_alloc,
+        sycl_queue=sycl_queue_alloc,
+    )
+
+    # Use np.copyto to burn each choicelist array onto result, using the
+    # corresponding condlist as a boolean mask. This is done in reverse
+    # order since the first choice should take precedence.
+    choicelist = choicelist[-2::-1]
+    condlist = condlist[::-1]
+    for choice, cond in zip(choicelist, condlist):
+        dpnp.copyto(result, choice, where=cond)
 
-    return call_origin(numpy.select, condlist, choicelist, default)
+    return result
 
 
 # pylint: disable=redefined-outer-name

tests/skipped_tests.tbl

@@ -116,19 +116,6 @@ tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_i
 tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_index
 tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_order
 
-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/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_complex
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_default
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_default_complex
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_default_scalar
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_empty_lists
-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_broadcastable_complex
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_non_broadcastable
-
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmaskDifferentDtypes::test_putmask_differnt_dtypes_raises
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmask::test_putmask_non_equal_shape_raises
 

tests/skipped_tests_gpu.tbl

@@ -169,20 +169,6 @@ tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_i
 tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_index
 tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_order
 
-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/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_complex
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_default
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_default_complex
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_default_scalar
-tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_empty_lists
-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_broadcastable_complex
-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_condlist
-
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmaskDifferentDtypes::test_putmask_differnt_dtypes_raises
 tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmask::test_putmask_non_equal_shape_raises
 

tests/test_sycl_queue.py

@@ -2338,6 +2338,19 @@ def test_astype(device_x, device_y):
     assert_sycl_queue_equal(y.sycl_queue, sycl_queue)
 
 
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_select(device):
+    sycl_queue = dpctl.SyclQueue(device)
+    condlist = [dpnp.array([True, False], sycl_queue=sycl_queue)]
+    choicelist = [dpnp.array([1, 2], sycl_queue=sycl_queue)]
+    res = dpnp.select(condlist, choicelist)
+    assert_sycl_queue_equal(res.sycl_queue, sycl_queue)
+
+
 @pytest.mark.parametrize("copy", [True, False], ids=["True", "False"])
 @pytest.mark.parametrize(
     "device",

tests/test_usm_type.py

@@ -1363,6 +1363,14 @@ def test_histogram_bin_edges(usm_type_v, usm_type_w):
     assert edges.usm_type == du.get_coerced_usm_type([usm_type_v, usm_type_w])
 
 
+@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
+def test_select(usm_type):
+    condlist = [dp.array([True, False], usm_type=usm_type)]
+    choicelist = [dp.array([1, 2], usm_type=usm_type)]
+    res = dp.select(condlist, choicelist)
+    assert res.usm_type == usm_type
+
+
 @pytest.mark.parametrize("copy", [True, False], ids=["True", "False"])
 @pytest.mark.parametrize("usm_type_a", list_of_usm_types, ids=list_of_usm_types)
 def test_nan_to_num(copy, usm_type_a):

tests/third_party/cupy/indexing_tests/test_indexing.py

@@ -355,10 +355,9 @@ def test_select_type_error_condlist(self, dtype):
         a = cupy.arange(10, dtype=dtype)
         condlist = [[3] * 10, [2] * 10]
         choicelist = [a, a**2]
-        with pytest.raises(AttributeError):
+        with pytest.raises(TypeError):
             cupy.select(condlist, choicelist)
 
-    @pytest.mark.usefixtures("allow_fall_back_on_numpy")
     @testing.for_all_dtypes(no_bool=True)
     def test_select_type_error_choicelist(self, dtype):
         a, b = list(range(10)), list(range(-10, 0))
@@ -388,7 +387,7 @@ def test_select_default_scalar(self, dtype):
         b = cupy.arange(20)
         condlist = [a < 3, b > 8]
         choicelist = [a, b]
-        with pytest.raises(TypeError):
+        with pytest.raises(ValueError):
             cupy.select(condlist, choicelist, [dtype(2)])
 
     @pytest.mark.skip("as_strided() is not implemented yet")