Implement dpnp.select

dpnp/dpnp_iface_indexing.py

@@ -1423,7 +1423,7 @@ def select(condlist, choicelist, default=0):
     dpnp.check_supported_arrays_type(*choicelist)
 
     if not dpnp.isscalar(default) and not (
-        dpnp.is_supported_array_type(default) and default.shape == ()
+        dpnp.is_supported_array_type(default) and default.ndim == 0
     ):
         raise TypeError(
             "A default value must be any of scalar or 0-d supported array type"
@@ -1456,9 +1456,7 @@ def select(condlist, choicelist, default=0):
         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.
+    # Do in reverse order since the first choice should take precedence.
     choicelist = choicelist[::-1]
     condlist = condlist[::-1]
     for choice, cond in zip(choicelist, condlist):

tests/test_indexing.py

@@ -13,7 +13,7 @@
 
 import dpnp
 
-from .helper import get_all_dtypes, get_integer_dtypes
+from .helper import get_all_dtypes, get_integer_dtypes, has_support_aspect64
 
 
 def _add_keepdims(func):
@@ -1057,3 +1057,87 @@ def test_fill_diagonal_error():
     arr = dpnp.ones((1, 2, 3))
     with pytest.raises(ValueError):
         dpnp.fill_diagonal(arr, 5)
+
+
+class TestSelect:
+    choices_np = [
+        numpy.array([1, 2, 3]),
+        numpy.array([4, 5, 6]),
+        numpy.array([7, 8, 9]),
+    ]
+    choices_dp = [
+        dpnp.array([1, 2, 3]),
+        dpnp.array([4, 5, 6]),
+        dpnp.array([7, 8, 9]),
+    ]
+    conditions_np = [
+        numpy.array([False, False, False]),
+        numpy.array([False, True, False]),
+        numpy.array([False, False, True]),
+    ]
+    conditions_dp = [
+        dpnp.array([False, False, False]),
+        dpnp.array([False, True, False]),
+        dpnp.array([False, False, True]),
+    ]
+
+    def test_basic(self):
+        expected = numpy.select(self.conditions_np, self.choices_np, default=15)
+        result = dpnp.select(self.conditions_dp, self.choices_dp, default=15)
+        assert_array_equal(expected, result)
+
+    def test_broadcasting(self):
+        conditions_np = [numpy.array(True), numpy.array([False, True, False])]
+        conditions_dp = [dpnp.array(True), dpnp.array([False, True, False])]
+        choices_np = [numpy.array(1), numpy.arange(12).reshape(4, 3)]
+        choices_dp = [dpnp.array(1), dpnp.arange(12).reshape(4, 3)]
+        expected = numpy.select(conditions_np, choices_np)
+        result = dpnp.select(conditions_dp, choices_dp)
+        assert_array_equal(expected, result)
+
+    def test_return_dtype(self):
+        dtype = dpnp.complex128 if has_support_aspect64() else dpnp.complex64
+        assert_equal(
+            dpnp.select(self.conditions_dp, self.choices_dp, 1j).dtype, dtype
+        )
+
+        choices = [choice.astype(dpnp.int32) for choice in self.choices_dp]
+        assert_equal(dpnp.select(self.conditions_dp, choices).dtype, dpnp.int32)
+
+    def test_nan(self):
+        choice_np = numpy.array([1, 2, 3, numpy.nan, 5, 7])
+        choice_dp = dpnp.array([1, 2, 3, dpnp.nan, 5, 7])
+        condition_np = numpy.isnan(choice_np)
+        condition_dp = dpnp.isnan(choice_dp)
+        expected = numpy.select([condition_np], [choice_np])
+        result = dpnp.select([condition_dp], [choice_dp])
+        assert_array_equal(expected, result)
+
+    def test_many_arguments(self):
+        condition_np = [numpy.array([False])] * 100
+        condition_dp = [dpnp.array([False])] * 100
+        choice_np = [numpy.array([1])] * 100
+        choice_dp = [dpnp.array([1])] * 100
+        expected = numpy.select(condition_np, choice_np)
+        result = dpnp.select(condition_dp, choice_dp)
+        assert_array_equal(expected, result)
+
+    def test_deprecated_empty(self):
+        assert_raises(ValueError, dpnp.select, [], [], 3j)
+        assert_raises(ValueError, dpnp.select, [], [])
+
+    def test_non_bool_deprecation(self):
+        choices = self.choices_dp
+        conditions = self.conditions_dp[:]
+        conditions[0] = conditions[0].astype(dpnp.int64)
+        assert_raises(TypeError, dpnp.select, conditions, choices)
+
+    def test_error(self):
+        x0 = dpnp.array([True, False])
+        x1 = dpnp.array([1, 2])
+        with pytest.raises(ValueError):
+            dpnp.select([x0], [x1, x1])
+        with pytest.raises(TypeError):
+            dpnp.select([x0], [x1], default=x1)
+        with pytest.raises(TypeError):
+            dpnp.select([x1], [x1])

tests/test_sycl_queue.py

@@ -2449,11 +2449,10 @@ def test_astype(device_x, device_y):
     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)]
+    condlist = [dpnp.array([True, False], device=device)]
+    choicelist = [dpnp.array([1, 2], device=device)]
     res = dpnp.select(condlist, choicelist)
-    assert_sycl_queue_equal(res.sycl_queue, sycl_queue)
+    assert_sycl_queue_equal(res.sycl_queue, condlist[0].sycl_queue)
 
 
 @pytest.mark.parametrize("axis", [None, 0, -1])

tests/test_usm_type.py

@@ -1416,12 +1416,13 @@ 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)]
+@pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("usm_type_y", list_of_usm_types, ids=list_of_usm_types)
+def test_select(usm_type_x, usm_type_y):
+    condlist = [dp.array([True, False], usm_type=usm_type_x)]
+    choicelist = [dp.array([1, 2], usm_type=usm_type_y)]
     res = dp.select(condlist, choicelist)
-    assert res.usm_type == usm_type
+    assert res.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_y])
 
 
 @pytest.mark.parametrize("axis", [None, 0, -1])