Adds more tests for choose

Author: ndgrigorian

dpnp/tests/test_indexing.py

@@ -5,6 +5,7 @@
 import numpy
 import pytest
 from dpctl.tensor._numpy_helper import AxisError
+from dpctl.tensor._type_utils import _to_device_supported_dtype
 from dpctl.utils import ExecutionPlacementError
 from numpy.testing import (
     assert_,
@@ -1507,3 +1508,103 @@ def test_choose_0d_inputs(self):
         chc = dpnp.ones(sh, dtype="i4")
         r = dpnp.choose(inds, [chc])
         assert r == chc
+
+    def test_choose_in_overlaps_out(self):
+        # overlap with inds
+        inds = dpnp.zeros(6, dtype="i4")
+        inds_np = dpnp.asnumpy(inds)
+        chc_np = numpy.arange(6, dtype="i4")
+        chc = dpnp.arange(6, dtype="i4")
+        out = inds
+        expected = numpy.choose(inds_np, chc_np)
+        result = dpnp.choose(inds, chc, out=out)
+        assert_array_equal(expected, result)
+        assert result is out
+        assert (inds == out).all()
+        # overlap with chc
+        inds = dpnp.zeros(6, dtype="i4")
+        out = chc
+        expected = numpy.choose(inds_np, chc_np)
+        result = dpnp.choose(inds, chc, out=out)
+        assert_array_equal(expected, result)
+        assert result is out
+        assert (inds == out).all()
+
+    def test_choose_strided(self):
+        # inds strided
+        inds = dpnp.tile(dpnp.array([0, 1], dtype="i4"), 5)
+        inds_np = dpnp.asnumpy(inds)
+        c1 = dpnp.arange(5, dtype="i4")
+        c2 = dpnp.full(5, -1, dtype="i4")
+        chcs = [c1, c2]
+        chcs_np = [dpnp.asnumpy(chc) for chc in chcs]
+        result = dpnp.choose(inds[::-2], chcs)
+        expected = numpy.choose(inds_np[::-2], chcs_np)
+        assert_array_equal(result, expected)
+        # choices strided
+        c3 = dpnp.arange(20, dtype="i4")
+        c4 = dpnp.full(20, -1, dtype="i4")
+        chcs = [c3[::-2], c4[::-2]]
+        chcs_np = [dpnp.asnumpy(c3)[::-2], dpnp.asnumpy(c4)[::-2]]
+        result = dpnp.choose(inds, chcs)
+        expected = numpy.choose(inds_np, chcs_np)
+        assert_array_equal(result, expected)
+        # all strided
+        result = dpnp.choose(inds[::-1], chcs)
+        expected = numpy.choose(inds_np[::-1], chcs_np)
+        assert_array_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "indices", [[0, 2], [-5, 4]], ids=["[0, 2]", "[-5, 4]"]
+    )
+    @pytest.mark.parametrize("mode", ["clip", "wrap"])
+    def test_choose_modes(self, indices, mode):
+        chc = dpnp.array([-2, -1, 0, 1, 2], dtype="i4")
+        chc_np = dpnp.asnumpy(chc)
+        inds = dpnp.array(indices, dtype="i4")
+        inds_np = dpnp.asnumpy(inds)
+        expected = numpy.choose(inds_np, chc_np, mode=mode)
+        result = dpnp.choose(inds, chc, mode=mode)
+        assert_array_equal(expected, result)
+
+    def test_choose_arg_validation(self):
+        with pytest.raises(TypeError):
+            dpnp.choose(dpnp.zeros(()), 1)
+        with pytest.raises(ValueError):
+            dpnp.choose(dpnp.zeros(()), dpnp.ones(()), mode="err")
+
+    # based on examples from NumPy
+    def test_choose_broadcasting(self):
+        inds = dpnp.array([[1, 0, 1], [0, 1, 0], [1, 0, 1]], dtype="i4")
+        inds_np = dpnp.asnumpy(inds)
+        chcs = dpnp.array([-10, 10])
+        chcs_np = dpnp.asnumpy(chcs)
+        result = dpnp.choose(inds, chcs)
+        expected = numpy.choose(inds_np, chcs_np)
+        assert_array_equal(result, expected)
+
+        inds = dpnp.array([0, 1]).reshape((2, 1, 1))
+        inds_np = dpnp.asnumpy(inds)
+        chc1 = dpnp.array([1, 2, 3]).reshape((1, 3, 1))
+        chc2 = dpnp.array([-1, -2, -3, -4, -5]).reshape(1, 1, 5)
+        chcs = [chc1, chc2]
+        chcs_np = [dpnp.asnumpy(chc) for chc in chcs]
+        result = dpnp.choose(inds, chcs)
+        expected = numpy.choose(inds_np, chcs_np)
+        assert_array_equal(result, expected)
+
+    @pytest.mark.parametrize("chc1_dt", get_all_dtypes())
+    @pytest.mark.parametrize("chc2_dt", get_all_dtypes())
+    def test_choose_promote_choices(self, chc1_dt, chc2_dt):
+        inds = dpnp.array([0, 1], dtype="i4")
+        inds_np = dpnp.asnumpy(inds)
+        chc1 = dpnp.zeros(1, dtype=chc1_dt)
+        chc2 = dpnp.ones(1, dtype=chc2_dt)
+        chcs = [chc1, chc2]
+        chcs_np = [dpnp.asnumpy(chc) for chc in chcs]
+        result = dpnp.choose(inds, chcs)
+        expected = numpy.choose(inds_np, chcs_np)
+        assert (
+            _to_device_supported_dtype(expected.dtype, inds.sycl_device)
+            == result.dtype
+        )

dpnp/tests/test_sycl_queue.py

@@ -2880,3 +2880,24 @@ def test_ix(device_0, device_1):
     ixgrid = dpnp.ix_(x0, x1)
     assert_sycl_queue_equal(ixgrid[0].sycl_queue, x0.sycl_queue)
     assert_sycl_queue_equal(ixgrid[1].sycl_queue, x1.sycl_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_choose(device):
+    chc = dpnp.arange(5, dtype="i4", device=device)
+    chc_np = dpnp.asnumpy(chc)
+
+    inds = dpnp.array([0, 1, 3], dtype="i4", device=device)
+    inds_np = dpnp.asnumpy(inds)
+
+    result = dpnp.choose(inds, chc)
+    expected = numpy.choose(inds_np, chc_np)
+    assert_allclose(expected, result)
+
+    expected_queue = chc.sycl_queue
+    result_queue = result.sycl_queue
+    assert_sycl_queue_equal(result_queue, expected_queue)

dpnp/tests/test_usm_type.py

@@ -1764,3 +1764,17 @@ def test_ix(usm_type_0, usm_type_1):
     ixgrid = dp.ix_(x0, x1)
     assert ixgrid[0].usm_type == x0.usm_type
     assert ixgrid[1].usm_type == x1.usm_type
+
+
+@pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize(
+    "usm_type_ind", list_of_usm_types, ids=list_of_usm_types
+)
+def test_choose(usm_type_x, usm_type_ind):
+    chc = dp.arange(5, usm_type=usm_type_x)
+    ind = dp.array([0, 2, 4], usm_type=usm_type_ind)
+    z = dp.choose(ind, chc)
+
+    assert chc.usm_type == usm_type_x
+    assert ind.usm_type == usm_type_ind
+    assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_ind])