Where changed for empty and F-contiguous input

- Where now outputs an F-contiguous array when all inputs are F-contiguous
- Where now outputs a empty 0D array if any input is a 0D empty array
- Added tests for these cases

Fixed incorrect logic in where test

Author: ndgrigorian

dpctl/tensor/_search_functions.py

@@ -72,14 +72,21 @@ def where(condition, x1, x2):
     x1_dtype = x1.dtype
     x2_dtype = x2.dtype
     dst_dtype = _where_result_type(x1_dtype, x2_dtype, exec_q.sycl_device)
-
-    if condition.size == 0:
-        return dpt.asarray(
-            (), dtype=dst_dtype, usm_type=dst_usm_type, sycl_queue=exec_q
+    if dst_dtype is None:
+        raise TypeError(
+            "function 'where' does not support input "
+            f"types ({x1_dtype}, {x2_dtype}), "
+            "and the inputs could not be safely coerced "
+            "to any supported types according to the casting rule ''safe''."
         )
 
     res_shape = _broadcast_shapes(condition, x1, x2)
 
+    if condition.size == 0:
+        return dpt.empty(
+            res_shape, dtype=dst_dtype, usm_type=dst_usm_type, sycl_queue=exec_q
+        )
+
     deps = []
     wait_list = []
     if x1_dtype is not dst_dtype:
@@ -104,8 +111,25 @@ def where(condition, x1, x2):
     x1 = dpt.broadcast_to(x1, res_shape)
     x2 = dpt.broadcast_to(x2, res_shape)
 
+    # dst is F-contiguous when all inputs are F contiguous
+    # otherwise, defaults to C-contiguous
+    if all(
+        (
+            condition.flags.fnc,
+            x1.flags.fnc,
+            x2.flags.fnc,
+        )
+    ):
+        order = "F"
+    else:
+        order = "C"
+
     dst = dpt.empty(
-        res_shape, dtype=dst_dtype, usm_type=dst_usm_type, sycl_queue=exec_q
+        res_shape,
+        dtype=dst_dtype,
+        order=order,
+        usm_type=dst_usm_type,
+        sycl_queue=exec_q,
     )
 
     hev, _ = ti._where(

dpctl/tests/test_usm_ndarray_search_functions.py

@@ -39,7 +39,7 @@
 
 
 def test_where_basic():
-    get_queue_or_skip
+    get_queue_or_skip()
 
     cond = dpt.asarray(
         [
@@ -58,27 +58,87 @@ def test_where_basic():
     assert (dpt.asnumpy(out) == dpt.asnumpy(out_expected)).all()
 
 
+def _dtype_all_close(x1, x2):
+    if np.issubdtype(x2.dtype, np.floating) or np.issubdtype(
+        x2.dtype, np.complexfloating
+    ):
+        x2_dtype = x2.dtype
+        return np.allclose(
+            x1, x2, atol=np.finfo(x2_dtype).eps, rtol=np.finfo(x2_dtype).eps
+        )
+    else:
+        return np.allclose(x1, x2)
+
+
 @pytest.mark.parametrize("dt1", _all_dtypes)
 @pytest.mark.parametrize("dt2", _all_dtypes)
 def test_where_all_dtypes(dt1, dt2):
     q = get_queue_or_skip()
     skip_if_dtype_not_supported(dt1, q)
     skip_if_dtype_not_supported(dt2, q)
 
-    cond_np = np.arange(5) > 2
-    x1_np = np.asarray(2, dtype=dt1)
-    x2_np = np.asarray(3, dtype=dt2)
-
-    cond = dpt.asarray(cond_np, sycl_queue=q)
-    x1 = dpt.asarray(x1_np, sycl_queue=q)
-    x2 = dpt.asarray(x2_np, sycl_queue=q)
+    cond = dpt.asarray([False, False, False, True, True], sycl_queue=q)
+    x1 = dpt.asarray(2, sycl_queue=q)
+    x2 = dpt.asarray(3, sycl_queue=q)
 
     res = dpt.where(cond, x1, x2)
-    res_np = np.where(cond_np, x1_np, x2_np)
+    res_check = np.asarray([3, 3, 3, 2, 2], dtype=res.dtype)
 
-    if res.dtype != res_np.dtype:
-        assert res.dtype.kind == res_np.dtype.kind
-        assert_array_equal(dpt.asnumpy(res).astype(res_np.dtype), res_np)
+    dev = q.sycl_device
 
-    else:
-        assert_array_equal(dpt.asnumpy(res), res_np)
+    if not dev.has_aspect_fp16 or not dev.has_aspect_fp64:
+        assert res.dtype.kind == dpt.result_type(x1.dtype, x2.dtype).kind
+
+    assert _dtype_all_close(dpt.asnumpy(res), res_check)
+
+
+def test_where_empty():
+    # check that numpy returns same results when
+    # handling empty arrays
+    get_queue_or_skip()
+
+    empty = dpt.empty(0)
+    m = dpt.asarray(True)
+    x1 = dpt.asarray(1)
+    x2 = dpt.asarray(2)
+    res = dpt.where(empty, x1, x2)
+
+    empty_np = np.empty(0)
+    m_np = dpt.asnumpy(m)
+    x1_np = dpt.asnumpy(x1)
+    x2_np = dpt.asnumpy(x2)
+    res_np = np.where(empty_np, x1_np, x2_np)
+
+    assert_array_equal(dpt.asnumpy(res), res_np)
+
+    res = dpt.where(m, empty, x2)
+    res_np = np.where(m_np, empty_np, x2_np)
+
+    assert_array_equal(dpt.asnumpy(res), res_np)
+
+
+@pytest.mark.parametrize("dt", _all_dtypes)
+@pytest.mark.parametrize("order", ["C", "F"])
+def test_where_contiguous(dt, order):
+    q = get_queue_or_skip()
+    skip_if_dtype_not_supported(dt, q)
+
+    cond = dpt.asarray(
+        [
+            [[True, False, False], [False, True, True]],
+            [[False, True, False], [True, False, True]],
+            [[False, False, True], [False, False, True]],
+            [[False, False, False], [True, False, True]],
+            [[True, True, True], [True, False, True]],
+        ],
+        sycl_queue=q,
+        order=order,
+    )
+
+    x1 = dpt.full(cond.shape, 2, dtype=dt, order=order, sycl_queue=q)
+    x2 = dpt.full(cond.shape, 3, dtype=dt, order=order, sycl_queue=q)
+
+    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)