Adjusted where test data types

- Set dtypes used in tests (where possible) to i4 to prevent allocation failures on GPUs without double support
- Made tests for strided where kernel data more robust

Author: ndgrigorian

dpctl/tensor/libtensor/source/where.cpp

@@ -166,8 +166,8 @@ py_where(dpctl::tensor::usm_ndarray condition,
 
         auto where_ev = contig_fn(exec_q, nelems, cond_data, x1_data, x2_data,
                                   dst_data, depends);
-        sycl::event ht_ev = dpctl::utils::keep_args_alive(
-            exec_q, {x1, x2, dst, condition}, {where_ev});
+        sycl::event ht_ev =
+            keep_args_alive(exec_q, {x1, x2, dst, condition}, {where_ev});
 
         return std::make_pair(ht_ev, where_ev);
     }

dpctl/tests/test_usm_ndarray_search_functions.py

@@ -70,8 +70,8 @@ def test_where_basic():
 
     out = dpt.where(
         cond,
-        dpt.ones(cond.shape[0])[:, dpt.newaxis],
-        dpt.zeros(cond.shape[0])[:, dpt.newaxis],
+        dpt.ones(cond.shape[0], dtype="i4")[:, dpt.newaxis],
+        dpt.zeros(cond.shape[0], dtype="i4")[:, dpt.newaxis],
     )
     assert (dpt.asnumpy(out) == dpt.asnumpy(out_expected)).all()
 
@@ -162,13 +162,13 @@ def test_where_empty():
     # handling empty arrays
     get_queue_or_skip()
 
-    empty = dpt.empty(0)
+    empty = dpt.empty(0, dtype="i2")
     m = dpt.asarray(True)
-    x1 = dpt.asarray(1)
-    x2 = dpt.asarray(2)
+    x1 = dpt.asarray(1, dtype="i2")
+    x2 = dpt.asarray(2, dtype="i2")
     res = dpt.where(empty, x1, x2)
 
-    empty_np = np.empty(0)
+    empty_np = np.empty(0, dtype="i2")
     m_np = dpt.asnumpy(m)
     x1_np = dpt.asnumpy(x1)
     x2_np = dpt.asnumpy(x2)
@@ -201,8 +201,8 @@ def test_where_contiguous(order):
         order=order,
     )
 
-    x1 = dpt.full(cond.shape, 2, order=order)
-    x2 = dpt.full(cond.shape, 3, order=order)
+    x1 = dpt.full(cond.shape, 2, dtype="i4", order=order)
+    x2 = dpt.full(cond.shape, 3, dtype="i4", order=order)
     expected = np.where(dpt.asnumpy(cond), dpt.asnumpy(x1), dpt.asnumpy(x2))
     res = dpt.where(cond, x1, x2)
 
@@ -214,11 +214,11 @@ def test_where_contiguous1D():
 
     cond = dpt.asarray([True, False, True, False, False, True])
 
-    x1 = dpt.full(cond.shape, 2)
-    x2 = dpt.full(cond.shape, 3)
+    x1 = dpt.full(cond.shape, 2, dtype="i4")
+    x2 = dpt.full(cond.shape, 3, dtype="i4")
     expected = np.where(dpt.asnumpy(cond), dpt.asnumpy(x1), dpt.asnumpy(x2))
     res = dpt.where(cond, x1, x2)
-    assert _dtype_all_close(dpt.asnumpy(res), expected)
+    assert_array_equal(dpt.asnumpy(res), expected)
 
     # test with complex dtype (branch in kernel)
     x1 = dpt.astype(x1, dpt.complex64)
@@ -239,20 +239,39 @@ def test_where_strided():
         (s0, s1),
     )[:, ::3]
 
-    x1 = dpt.ones((cond.shape[0], cond.shape[1] * 2))[:, ::2]
-    x2 = dpt.zeros((cond.shape[0], cond.shape[1] * 3))[:, ::3]
+    x1 = dpt.reshape(
+        dpt.arange(cond.shape[0] * cond.shape[1] * 2, dtype="i4"),
+        (cond.shape[0], cond.shape[1] * 2),
+    )[:, ::2]
+    x2 = dpt.reshape(
+        dpt.arange(cond.shape[0] * cond.shape[1] * 3, dtype="i4"),
+        (cond.shape[0], cond.shape[1] * 3),
+    )[:, ::3]
     expected = np.where(dpt.asnumpy(cond), dpt.asnumpy(x1), dpt.asnumpy(x2))
     res = dpt.where(cond, x1, x2)
 
-    assert _dtype_all_close(dpt.asnumpy(res), expected)
+    assert_array_equal(dpt.asnumpy(res), expected)
+
+    # negative strides
+    res = dpt.where(cond, dpt.flip(x1), x2)
+    expected = np.where(
+        dpt.asnumpy(cond), np.flip(dpt.asnumpy(x1)), dpt.asnumpy(x2)
+    )
+    assert_array_equal(dpt.asnumpy(res), expected)
+
+    res = dpt.where(dpt.flip(cond), x1, x2)
+    expected = np.where(
+        np.flip(dpt.asnumpy(cond)), dpt.asnumpy(x1), dpt.asnumpy(x2)
+    )
+    assert_array_equal(dpt.asnumpy(res), expected)
 
 
 def test_where_arg_validation():
     get_queue_or_skip()
 
     check = dict()
-    x1 = dpt.empty((1,))
-    x2 = dpt.empty((1,))
+    x1 = dpt.empty((1,), dtype="i4")
+    x2 = dpt.empty((1,), dtype="i4")
 
     with pytest.raises(TypeError):
         dpt.where(check, x1, x2)
@@ -267,12 +286,12 @@ def test_where_compute_follows_data():
     q2 = get_queue_or_skip()
     q3 = get_queue_or_skip()
 
-    x1 = dpt.empty((1,), sycl_queue=q1)
-    x2 = dpt.empty((1,), sycl_queue=q2)
+    x1 = dpt.empty((1,), dtype="i4", sycl_queue=q1)
+    x2 = dpt.empty((1,), dtype="i4", sycl_queue=q2)
 
     with pytest.raises(ExecutionPlacementError):
-        dpt.where(dpt.empty((1,), sycl_queue=q1), x1, x2)
+        dpt.where(dpt.empty((1,), dtype="i4", sycl_queue=q1), x1, x2)
     with pytest.raises(ExecutionPlacementError):
-        dpt.where(dpt.empty((1,), sycl_queue=q3), x1, x2)
+        dpt.where(dpt.empty((1,), dtype="i4", sycl_queue=q3), x1, x2)
     with pytest.raises(ExecutionPlacementError):
         dpt.where(x1, x1, x2)