Added tests for where and type utility functions

Author: ndgrigorian

dpctl/tests/test_type_utils.py

@@ -0,0 +1,68 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2023 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+import dpctl.tensor as dpt
+from dpctl.tensor._type_utils import (
+    _all_data_types,
+    _can_cast,
+    _is_maximal_inexact_type,
+)
+
+
+def test_all_data_types():
+    fp16_fp64_types = set([dpt.float16, dpt.float64, dpt.complex128])
+    fp64_types = set([dpt.float64, dpt.complex128])
+
+    all_dts = _all_data_types(True, True)
+    assert fp16_fp64_types.issubset(all_dts)
+
+    all_dts = _all_data_types(True, False)
+    assert dpt.float16 in all_dts
+    assert not fp64_types.issubset(all_dts)
+
+    all_dts = _all_data_types(False, True)
+    assert dpt.float16 not in all_dts
+    assert fp64_types.issubset(all_dts)
+
+    all_dts = _all_data_types(False, False)
+    assert not fp16_fp64_types.issubset(all_dts)
+
+
+@pytest.mark.parametrize("fp16", [True, False])
+@pytest.mark.parametrize("fp64", [True, False])
+def test_maximal_inexact_types(fp16, fp64):
+    assert not _is_maximal_inexact_type(dpt.int32, fp16, fp64)
+    assert fp64 == _is_maximal_inexact_type(dpt.float64, fp16, fp64)
+    assert fp64 == _is_maximal_inexact_type(dpt.complex128, fp16, fp64)
+    assert fp64 != _is_maximal_inexact_type(dpt.float32, fp16, fp64)
+    assert fp64 != _is_maximal_inexact_type(dpt.complex64, fp16, fp64)
+
+
+def test_can_cast_device():
+    assert _can_cast(dpt.int64, dpt.float64, True, True)
+    # if f8 is available, can't cast i8 to f4
+    assert not _can_cast(dpt.int64, dpt.float32, True, True)
+    assert not _can_cast(dpt.int64, dpt.float32, False, True)
+    # should be able to cast to f8 when f2 unavailable
+    assert _can_cast(dpt.int64, dpt.float64, False, True)
+    # casting to f4 acceptable when f8 unavailable
+    assert _can_cast(dpt.int64, dpt.float32, True, False)
+    assert _can_cast(dpt.int64, dpt.float32, False, False)
+    # can't safely cast inexact type to inexact type of lesser precision
+    assert not _can_cast(dpt.float32, dpt.float16, True, False)
+    assert not _can_cast(dpt.float64, dpt.float32, False, True)

dpctl/tests/test_usm_ndarray_search_functions.py

@@ -1,6 +1,6 @@
 #                      Data Parallel Control (dpctl)
 #
-# Copyright 2020-2022 Intel Corporation
+# Copyright 2020-2023 Intel Corporation
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,8 +20,12 @@
 from numpy.testing import assert_array_equal
 
 import dpctl.tensor as dpt
+from dpctl.tensor._search_functions import _where_result_type
+from dpctl.tensor._type_utils import _all_data_types
+from dpctl.utils import ExecutionPlacementError
 
 _all_dtypes = [
+    "?",
     "u1",
     "i1",
     "u2",
@@ -38,6 +42,12 @@
 ]
 
 
+class mock_device:
+    def __init__(self, fp16, fp64):
+        self.has_aspect_fp16 = fp16
+        self.has_aspect_fp64 = fp64
+
+
 def test_where_basic():
     get_queue_or_skip()
 
@@ -54,7 +64,16 @@ def test_where_basic():
     out_expected = dpt.asarray(
         [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 0], [1, 1, 1]]
     )
+    assert (dpt.asnumpy(out) == dpt.asnumpy(out_expected)).all()
 
+    out = dpt.where(cond, dpt.ones(cond.shape), dpt.zeros(cond.shape))
+    assert (dpt.asnumpy(out) == dpt.asnumpy(out_expected)).all()
+
+    out = dpt.where(
+        cond,
+        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()
 
 
@@ -72,38 +91,98 @@ def _dtype_all_close(x1, x2):
 
 @pytest.mark.parametrize("dt1", _all_dtypes)
 @pytest.mark.parametrize("dt2", _all_dtypes)
-def test_where_all_dtypes(dt1, dt2):
+@pytest.mark.parametrize("fp16", [True, False])
+@pytest.mark.parametrize("fp64", [True, False])
+def test_where_result_types(dt1, dt2, fp16, fp64):
+    dev = mock_device(fp16, fp64)
+
+    dt1 = dpt.dtype(dt1)
+    dt2 = dpt.dtype(dt2)
+    res_t = _where_result_type(dt1, dt2, dev)
+
+    if fp16 and fp64:
+        assert res_t == dpt.result_type(dt1, dt2)
+    else:
+        if res_t:
+            assert res_t.kind == dpt.result_type(dt1, dt2).kind
+        else:
+            # some illegal cases are covered above, but
+            # this guarantees that _where_result_type
+            # produces None only when one of the dtypes
+            # is illegal given fp aspects of device
+            all_dts = _all_data_types(fp16, fp64)
+            assert dt1 not in all_dts or dt2 not in all_dts
+
+
+@pytest.mark.parametrize("dt", _all_dtypes)
+def test_where_all_dtypes(dt):
     q = get_queue_or_skip()
-    skip_if_dtype_not_supported(dt1, q)
-    skip_if_dtype_not_supported(dt2, q)
+    skip_if_dtype_not_supported(dt, 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)
+    # mask dtype changes
+    cond = dpt.asarray([0, 1, 3, 0, 10], dtype=dt, sycl_queue=q)
+    x1 = dpt.asarray(0, dtype="f", sycl_queue=q)
+    x2 = dpt.asarray(1, dtype="f", sycl_queue=q)
+    res = dpt.where(cond, x1, x2)
+
+    res_check = np.asarray([1, 0, 0, 1, 0], dtype=res.dtype)
+    assert _dtype_all_close(dpt.asnumpy(res), res_check)
 
+    # contiguous cases
+    x1 = dpt.full(cond.shape, 0, dtype="f4", sycl_queue=q)
+    x2 = dpt.full(cond.shape, 1, dtype="f4", sycl_queue=q)
     res = dpt.where(cond, x1, x2)
-    res_check = np.asarray([3, 3, 3, 2, 2], dtype=res.dtype)
+    assert _dtype_all_close(dpt.asnumpy(res), res_check)
 
-    dev = q.sycl_device
+    # input array dtype changes
+    cond = dpt.asarray([False, True, True, False, True], sycl_queue=q)
+    x1 = dpt.asarray(0, dtype=dt, sycl_queue=q)
+    x2 = dpt.asarray(1, dtype=dt, sycl_queue=q)
+    res = dpt.where(cond, x1, x2)
 
-    if not dev.has_aspect_fp16 or not dev.has_aspect_fp64:
-        assert res.dtype.kind == dpt.result_type(x1.dtype, x2.dtype).kind
+    res_check = np.asarray([1, 0, 0, 1, 0], dtype=res.dtype)
+    assert _dtype_all_close(dpt.asnumpy(res), res_check)
 
+    # contiguous cases
+    x1 = dpt.full(cond.shape, 0, dtype=dt, sycl_queue=q)
+    x2 = dpt.full(cond.shape, 1, dtype=dt, sycl_queue=q)
+    res = dpt.where(cond, x1, x2)
     assert _dtype_all_close(dpt.asnumpy(res), res_check)
 
 
+def test_where_nan_inf():
+    get_queue_or_skip()
+
+    cond = dpt.asarray([True, False, True, False], dtype="?")
+    x1 = dpt.asarray([np.nan, 2.0, np.inf, 3.0], dtype="f4")
+    x2 = dpt.asarray([2.0, np.nan, 3.0, np.inf], dtype="f4")
+
+    cond_np = dpt.asnumpy(cond)
+    x1_np = dpt.asnumpy(x1)
+    x2_np = dpt.asnumpy(x2)
+
+    res = dpt.where(cond, x1, x2)
+    res_np = np.where(cond_np, x1_np, x2_np)
+
+    assert np.allclose(dpt.asnumpy(res), res_np, equal_nan=True)
+
+    res = dpt.where(x1, cond, x2)
+    res_np = np.where(x1_np, cond_np, x2_np)
+    assert _dtype_all_close(dpt.asnumpy(res), res_np)
+
+
 def test_where_empty():
     # check that numpy returns same results when
     # 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)
@@ -116,12 +195,14 @@ def test_where_empty():
 
     assert_array_equal(dpt.asnumpy(res), res_np)
 
+    # check that broadcasting is performed
+    with pytest.raises(ValueError):
+        dpt.where(empty, x1, dpt.empty((1, 2)))
+
 
-@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)
+def test_where_contiguous(order):
+    get_queue_or_skip()
 
     cond = dpt.asarray(
         [
@@ -131,14 +212,100 @@ def test_where_contiguous(dt, order):
             [[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)
+    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)
+
+    assert _dtype_all_close(dpt.asnumpy(res), expected)
+
+
+def test_where_contiguous1D():
+    get_queue_or_skip()
 
+    cond = dpt.asarray([True, False, True, False, False, True])
+
+    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_array_equal(dpt.asnumpy(res), expected)
 
+    # test with complex dtype (branch in kernel)
+    x1 = dpt.astype(x1, dpt.complex64)
+    x2 = dpt.astype(x2, dpt.complex64)
+    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)
+
+
+def test_where_strided():
+    get_queue_or_skip()
+
+    s0, s1 = 4, 9
+    cond = dpt.reshape(
+        dpt.asarray(
+            [True, False, False, False, True, True, False, True, False] * s0
+        ),
+        (s0, s1),
+    )[:, ::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_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,), dtype="i4")
+    x2 = dpt.empty((1,), dtype="i4")
+
+    with pytest.raises(TypeError):
+        dpt.where(check, x1, x2)
+    with pytest.raises(TypeError):
+        dpt.where(x1, check, x2)
+    with pytest.raises(TypeError):
+        dpt.where(x1, x2, check)
+
+
+def test_where_compute_follows_data():
+    q1 = get_queue_or_skip()
+    q2 = get_queue_or_skip()
+    q3 = get_queue_or_skip()
+
+    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,), dtype="i4", sycl_queue=q1), x1, x2)
+    with pytest.raises(ExecutionPlacementError):
+        dpt.where(dpt.empty((1,), dtype="i4", sycl_queue=q3), x1, x2)
+    with pytest.raises(ExecutionPlacementError):
+        dpt.where(x1, x1, x2)