11468253350

Committed 22 Oct 2024 08:46PM UTC coverage: 87.681% (+0.008%) from 87.673%

Build # 11468253350

Build Type

Pull #1876

github

Committed by

web-flow

Commit Message

Merge 0c6970695 into 9eb8f0303

Pull Request Pull Request #1876: result_type to raise unless at least one argument is usm_ndarray

Run Details

3077 of 3590 branches covered (85.71%)

Branch coverage included in aggregate %.

12 of 12 new or added lines in 2 files covered. (100.0%)

4 existing lines in 1 file now uncovered.

11720 of 13286 relevant lines covered (88.21%)

7160.11 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.65

/dpctl/tensor/_search_functions.py

#                       Data Parallel Control (dpctl)
#
#  Copyright 2020-2024 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 dpctl
import dpctl.tensor as dpt
import dpctl.tensor._tensor_impl as ti
from dpctl.tensor._elementwise_common import (
    _get_dtype,
    _get_queue_usm_type,
    _get_shape,
    _validate_dtype,
)
from dpctl.tensor._manipulation_functions import _broadcast_shape_impl
from dpctl.utils import ExecutionPlacementError, SequentialOrderManager

from ._copy_utils import _empty_like_orderK, _empty_like_triple_orderK
from ._type_utils import (
    WeakBooleanType,
    WeakComplexType,
    WeakFloatingType,
    WeakIntegralType,
    _all_data_types,
    _can_cast,
    _is_weak_dtype,
    _result_type_fn_impl,
    _strong_dtype_num_kind,
    _to_device_supported_dtype,
    _weak_type_num_kind,
)


def _default_dtype_from_weak_type(dt, dev):
    if isinstance(dt, WeakBooleanType):
        return dpt.bool
    if isinstance(dt, WeakIntegralType):
        return dpt.dtype(ti.default_device_int_type(dev))
    if isinstance(dt, WeakFloatingType):
        return dpt.dtype(ti.default_device_fp_type(dev))
    if isinstance(dt, WeakComplexType):
        return dpt.dtype(ti.default_device_complex_type(dev))


def _resolve_two_weak_types(o1_dtype, o2_dtype, dev):
    "Resolves two weak data types per NEP-0050"
    if _is_weak_dtype(o1_dtype):
        if _is_weak_dtype(o2_dtype):
            return _default_dtype_from_weak_type(
                o1_dtype, dev
            ), _default_dtype_from_weak_type(o2_dtype, dev)
        o1_kind_num = _weak_type_num_kind(o1_dtype)
        o2_kind_num = _strong_dtype_num_kind(o2_dtype)
        if o1_kind_num > o2_kind_num:
            if isinstance(o1_dtype, WeakIntegralType):
                return dpt.dtype(ti.default_device_int_type(dev)), o2_dtype
            if isinstance(o1_dtype, WeakComplexType):
                if o2_dtype is dpt.float16 or o2_dtype is dpt.float32:
                    return dpt.complex64, o2_dtype
                return (
                    _to_device_supported_dtype(dpt.complex128, dev),
                    o2_dtype,
                )
            return _to_device_supported_dtype(dpt.float64, dev), o2_dtype
        else:
            return o2_dtype, o2_dtype
    elif _is_weak_dtype(o2_dtype):
        o1_kind_num = _strong_dtype_num_kind(o1_dtype)
        o2_kind_num = _weak_type_num_kind(o2_dtype)
        if o2_kind_num > o1_kind_num:
            if isinstance(o2_dtype, WeakIntegralType):
                return o1_dtype, dpt.dtype(ti.default_device_int_type(dev))
            if isinstance(o2_dtype, WeakComplexType):
                if o1_dtype is dpt.float16 or o1_dtype is dpt.float32:
                    return o1_dtype, dpt.complex64
                return o1_dtype, _to_device_supported_dtype(dpt.complex128, dev)
            return (
                o1_dtype,
                _to_device_supported_dtype(dpt.float64, dev),
            )
        else:
            return o1_dtype, o1_dtype
    else:
        return o1_dtype, o2_dtype


def _where_result_type(dt1, dt2, dev):
    res_dtype = _result_type_fn_impl(
        (
            dt1,
            dt2,
        ),
        sycl_device=dev,
    )
    fp16 = dev.has_aspect_fp16
    fp64 = dev.has_aspect_fp64

    all_dts = _all_data_types(fp16, fp64)
    if res_dtype in all_dts:
        return res_dtype
    else:
        for res_dtype_ in all_dts:
            if _can_cast(dt1, res_dtype_, fp16, fp64) and _can_cast(
                dt2, res_dtype_, fp16, fp64
            ):
                return res_dtype_
        return None


def where(condition, x1, x2, /, *, order="K", out=None):
    """
    Returns :class:`dpctl.tensor.usm_ndarray` with elements chosen
    from ``x1`` or ``x2`` depending on ``condition``.

    Args:
        condition (usm_ndarray): When ``True`` yields from ``x1``,
            and otherwise yields from ``x2``.
            Must be compatible with ``x1`` and ``x2`` according
            to broadcasting rules.
        x1 (Union[usm_ndarray, bool, int, float, complex]):
            Array from which values are chosen when ``condition`` is ``True``.
            Must be compatible with ``condition`` and ``x2`` according
            to broadcasting rules.
        x2 (Union[usm_ndarray, bool, int, float, complex]):
            Array from which values are chosen when ``condition`` is not
            ``True``.
            Must be compatible with ``condition`` and ``x2`` according
            to broadcasting rules.
        order (``"K"``, ``"C"``, ``"F"``, ``"A"``, optional):
            Memory layout of the new output array,
            if parameter ``out`` is ``None``.
            Default: ``"K"``.
        out (Optional[usm_ndarray]):
            the array into which the result is written.
            The data type of `out` must match the expected shape and the
            expected data type of the result.
            If ``None`` then a new array is returned. Default: ``None``.

    Returns:
        usm_ndarray:
            An array with elements from ``x1`` where ``condition`` is ``True``,
            and elements from ``x2`` elsewhere.

    The data type of the returned array is determined by applying
    the Type Promotion Rules to ``x1`` and ``x2``.
    """
    if not isinstance(condition, dpt.usm_ndarray):
        raise TypeError(
            "Expecting dpctl.tensor.usm_ndarray type, " f"got {type(condition)}"
        )
    if order not in ["K", "C", "F", "A"]:
        order = "K"
    q1, condition_usm_type = condition.sycl_queue, condition.usm_type
    q2, x1_usm_type = _get_queue_usm_type(x1)
    q3, x2_usm_type = _get_queue_usm_type(x2)
    if q2 is None and q3 is None:
        exec_q = q1
        out_usm_type = condition_usm_type
    elif q3 is None:
        exec_q = dpctl.utils.get_execution_queue((q1, q2))
        if exec_q is None:
            raise ExecutionPlacementError(
                "Execution placement can not be unambiguously inferred "
                "from input arguments."
            )
        out_usm_type = dpctl.utils.get_coerced_usm_type(
            (
                condition_usm_type,
                x1_usm_type,
            )
        )
    elif q2 is None:
        exec_q = dpctl.utils.get_execution_queue((q1, q3))
        if exec_q is None:
            raise ExecutionPlacementError(
                "Execution placement can not be unambiguously inferred "
                "from input arguments."
            )
        out_usm_type = dpctl.utils.get_coerced_usm_type(
            (
                condition_usm_type,
                x2_usm_type,
            )
        )
    else:
        exec_q = dpctl.utils.get_execution_queue((q1, q2, q3))
        if exec_q is None:
            raise ExecutionPlacementError(
                "Execution placement can not be unambiguously inferred "
                "from input arguments."
            )
        out_usm_type = dpctl.utils.get_coerced_usm_type(
            (
                condition_usm_type,
                x1_usm_type,
                x2_usm_type,
            )
        )
    dpctl.utils.validate_usm_type(out_usm_type, allow_none=False)
    condition_shape = condition.shape
    x1_shape = _get_shape(x1)
    x2_shape = _get_shape(x2)
    if not all(
        isinstance(s, (tuple, list))
        for s in (
            x1_shape,
            x2_shape,
        )
    ):
        raise TypeError(
            "Shape of arguments can not be inferred. "
            "Arguments are expected to be "
            "lists, tuples, or both"
        )
    try:
        res_shape = _broadcast_shape_impl(
            [
                condition_shape,
                x1_shape,
                x2_shape,
            ]
        )
    except ValueError:
        raise ValueError(
            "operands could not be broadcast together with shapes "
            f"{condition_shape}, {x1_shape}, and {x2_shape}"
        )
    sycl_dev = exec_q.sycl_device
    x1_dtype = _get_dtype(x1, sycl_dev)
    x2_dtype = _get_dtype(x2, sycl_dev)
    if not all(_validate_dtype(o) for o in (x1_dtype, x2_dtype)):
        raise ValueError("Operands have unsupported data types")
    x1_dtype, x2_dtype = _resolve_two_weak_types(x1_dtype, x2_dtype, sycl_dev)
    out_dtype = _where_result_type(x1_dtype, x2_dtype, sycl_dev)
    if out_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''."
        )

    orig_out = out
    if out is not None:
        if not isinstance(out, dpt.usm_ndarray):
            raise TypeError(
                "output array must be of usm_ndarray type, got " f"{type(out)}"
            )

        if not out.flags.writable:
            raise ValueError("provided `out` array is read-only")

        if out.shape != res_shape:
            raise ValueError(
                "The shape of input and output arrays are "
                f"inconsistent. Expected output shape is {res_shape}, "
                f"got {out.shape}"
            )

        if out_dtype != out.dtype:
            raise ValueError(
                f"Output array of type {out_dtype} is needed, "
                f"got {out.dtype}"
            )

        if dpctl.utils.get_execution_queue((exec_q, out.sycl_queue)) is None:
            raise ExecutionPlacementError(
                "Input and output allocation queues are not compatible"
            )

        if ti._array_overlap(condition, out) and not ti._same_logical_tensors(
            condition, out
        ):
            out = dpt.empty_like(out)

        if isinstance(x1, dpt.usm_ndarray):
            if (
                ti._array_overlap(x1, out)
                and not ti._same_logical_tensors(x1, out)
                and x1_dtype == out_dtype
            ):
                out = dpt.empty_like(out)

        if isinstance(x2, dpt.usm_ndarray):
            if (
                ti._array_overlap(x2, out)
                and not ti._same_logical_tensors(x2, out)
                and x2_dtype == out_dtype
            ):
                out = dpt.empty_like(out)

    if order == "A":
        order = (
            "F"
            if all(
                arr.flags.f_contiguous
                for arr in (
                    condition,
                    x1,
                    x2,
                )
            )
            else "C"
        )
    if not isinstance(x1, dpt.usm_ndarray):
        x1 = dpt.asarray(x1, dtype=x1_dtype, sycl_queue=exec_q)
    if not isinstance(x2, dpt.usm_ndarray):
        x2 = dpt.asarray(x2, dtype=x2_dtype, sycl_queue=exec_q)

    if condition.size == 0:
        if out is not None:
            return out
        else:
            if order == "K":
                return _empty_like_triple_orderK(
                    condition,
                    x1,
                    x2,
                    out_dtype,
                    res_shape,
                    out_usm_type,
                    exec_q,
                )
            else:
                return dpt.empty(
                    res_shape,
                    dtype=out_dtype,
                    order=order,
                    usm_type=out_usm_type,
                    sycl_queue=exec_q,
                )

    _manager = SequentialOrderManager[exec_q]
    dep_evs = _manager.submitted_events
    if x1_dtype != out_dtype:
        if order == "K":
            _x1 = _empty_like_orderK(x1, out_dtype)
        else:
            _x1 = dpt.empty_like(x1, dtype=out_dtype, order=order)
        ht_copy1_ev, copy1_ev = ti._copy_usm_ndarray_into_usm_ndarray(
            src=x1, dst=_x1, sycl_queue=exec_q, depends=dep_evs
        )
        x1 = _x1
        _manager.add_event_pair(ht_copy1_ev, copy1_ev)

    if x2_dtype != out_dtype:
        if order == "K":
            _x2 = _empty_like_orderK(x2, out_dtype)
        else:
            _x2 = dpt.empty_like(x2, dtype=out_dtype, order=order)
        ht_copy2_ev, copy2_ev = ti._copy_usm_ndarray_into_usm_ndarray(
            src=x2, dst=_x2, sycl_queue=exec_q, depends=dep_evs
        )
        x2 = _x2
        _manager.add_event_pair(ht_copy2_ev, copy2_ev)

    if out is None:
        if order == "K":
            out = _empty_like_triple_orderK(
                condition, x1, x2, out_dtype, res_shape, out_usm_type, exec_q
            )
        else:
            out = dpt.empty(
                res_shape,
                dtype=out_dtype,
                order=order,
                usm_type=out_usm_type,
                sycl_queue=exec_q,
            )

    if condition_shape != res_shape:
        condition = dpt.broadcast_to(condition, res_shape)
    if x1_shape != res_shape:
        x1 = dpt.broadcast_to(x1, res_shape)
    if x2_shape != res_shape:
        x2 = dpt.broadcast_to(x2, res_shape)

    dep_evs = _manager.submitted_events
    hev, where_ev = ti._where(
        condition=condition,
        x1=x1,
        x2=x2,
        dst=out,
        sycl_queue=exec_q,
        depends=dep_evs,
    )
    _manager.add_event_pair(hev, where_ev)
    if not (orig_out is None or orig_out is out):
        # Copy the out data from temporary buffer to original memory
        ht_copy_out_ev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
            src=out,
            dst=orig_out,
            sycl_queue=exec_q,
            depends=[where_ev],
        )
        _manager.add_event_pair(ht_copy_out_ev, cpy_ev)
        out = orig_out

    return out

1	# Data Parallel Control (dpctl)
2	#
3	# Copyright 2020-2024 Intel Corporation
4	#
5	# Licensed under the Apache License, Version 2.0 (the "License");
6	# you may not use this file except in compliance with the License.
7	# You may obtain a copy of the License at
8	#
9	# http://www.apache.org/licenses/LICENSE-2.0
10	#
11	# Unless required by applicable law or agreed to in writing, software
12	# distributed under the License is distributed on an "AS IS" BASIS,
13	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14	# See the License for the specific language governing permissions and
15	# limitations under the License.
16
17	import dpctl	1✔
18	import dpctl.tensor as dpt	1✔
19	import dpctl.tensor._tensor_impl as ti	1✔
20	from dpctl.tensor._elementwise_common import (	1✔
21	_get_dtype,
22	_get_queue_usm_type,
23	_get_shape,
24	_validate_dtype,
25	)
26	from dpctl.tensor._manipulation_functions import _broadcast_shape_impl	1✔
27	from dpctl.utils import ExecutionPlacementError, SequentialOrderManager	1✔
28
29	from ._copy_utils import _empty_like_orderK, _empty_like_triple_orderK	1✔
30	from ._type_utils import (	1✔
31	WeakBooleanType,
32	WeakComplexType,
33	WeakFloatingType,
34	WeakIntegralType,
35	_all_data_types,
36	_can_cast,
37	_is_weak_dtype,
38	_result_type_fn_impl,
39	_strong_dtype_num_kind,
40	_to_device_supported_dtype,
41	_weak_type_num_kind,
42	)
43
44
45	def _default_dtype_from_weak_type(dt, dev):	1✔
46	if isinstance(dt, WeakBooleanType):	1✔
47	return dpt.bool	1✔
48	if isinstance(dt, WeakIntegralType):	1✔
49	return dpt.dtype(ti.default_device_int_type(dev))	1✔
50	if isinstance(dt, WeakFloatingType):	1✔
51	return dpt.dtype(ti.default_device_fp_type(dev))	1✔
52	if isinstance(dt, WeakComplexType):	1!
53	return dpt.dtype(ti.default_device_complex_type(dev))	1✔
54
55
56	def _resolve_two_weak_types(o1_dtype, o2_dtype, dev):	1✔
57	"Resolves two weak data types per NEP-0050"
58	if _is_weak_dtype(o1_dtype):	1✔
59	if _is_weak_dtype(o2_dtype):	1✔
60	return _default_dtype_from_weak_type(	1✔
61	o1_dtype, dev
62	), _default_dtype_from_weak_type(o2_dtype, dev)
63	o1_kind_num = _weak_type_num_kind(o1_dtype)	1✔
64	o2_kind_num = _strong_dtype_num_kind(o2_dtype)	1✔
65	if o1_kind_num > o2_kind_num:	1✔
66	if isinstance(o1_dtype, WeakIntegralType):	1✔
67	return dpt.dtype(ti.default_device_int_type(dev)), o2_dtype	1✔
68	if isinstance(o1_dtype, WeakComplexType):	1✔
69	if o2_dtype is dpt.float16 or o2_dtype is dpt.float32:	1✔
70	return dpt.complex64, o2_dtype	1✔
71	return (	1✔
72	_to_device_supported_dtype(dpt.complex128, dev),
73	o2_dtype,
74	)
75	return _to_device_supported_dtype(dpt.float64, dev), o2_dtype	1✔
76	else:
77	return o2_dtype, o2_dtype	1✔
78	elif _is_weak_dtype(o2_dtype):	1✔
79	o1_kind_num = _strong_dtype_num_kind(o1_dtype)	1✔
80	o2_kind_num = _weak_type_num_kind(o2_dtype)	1✔
81	if o2_kind_num > o1_kind_num:	1✔
82	if isinstance(o2_dtype, WeakIntegralType):	1✔
83	return o1_dtype, dpt.dtype(ti.default_device_int_type(dev))	1✔
84	if isinstance(o2_dtype, WeakComplexType):	1✔
85	if o1_dtype is dpt.float16 or o1_dtype is dpt.float32:	1✔
86	return o1_dtype, dpt.complex64	1✔
87	return o1_dtype, _to_device_supported_dtype(dpt.complex128, dev)	1✔
88	return (	1✔
89	o1_dtype,
90	_to_device_supported_dtype(dpt.float64, dev),
91	)
92	else:
93	return o1_dtype, o1_dtype	1✔
94	else:
95	return o1_dtype, o2_dtype	1✔
96
97
98	def _where_result_type(dt1, dt2, dev):	1✔
99	res_dtype = _result_type_fn_impl(	1✔
100	(
101	dt1,
102	dt2,
103	),
104	sycl_device=dev,
105	)
106	fp16 = dev.has_aspect_fp16	1✔
107	fp64 = dev.has_aspect_fp64	1✔
108
109	all_dts = _all_data_types(fp16, fp64)	1✔
110	if res_dtype in all_dts:	1!
111	return res_dtype	1✔
112	else:
UNCOV 113	for res_dtype_ in all_dts:	×
UNCOV 114	if _can_cast(dt1, res_dtype_, fp16, fp64) and _can_cast(	×
115	dt2, res_dtype_, fp16, fp64
116	):
UNCOV 117	return res_dtype_	×
UNCOV 118	return None	×
119
120
121	def where(condition, x1, x2, /, *, order="K", out=None):	1✔
122	"""
123	Returns :class:`dpctl.tensor.usm_ndarray` with elements chosen
124	from ``x1`` or ``x2`` depending on ``condition``.
125
126	Args:
127	condition (usm_ndarray): When ``True`` yields from ``x1``,
128	and otherwise yields from ``x2``.
129	Must be compatible with ``x1`` and ``x2`` according
130	to broadcasting rules.
131	x1 (Union[usm_ndarray, bool, int, float, complex]):
132	Array from which values are chosen when ``condition`` is ``True``.
133	Must be compatible with ``condition`` and ``x2`` according
134	to broadcasting rules.
135	x2 (Union[usm_ndarray, bool, int, float, complex]):
136	Array from which values are chosen when ``condition`` is not
137	``True``.
138	Must be compatible with ``condition`` and ``x2`` according
139	to broadcasting rules.
140	order (``"K"``, ``"C"``, ``"F"``, ``"A"``, optional):
141	Memory layout of the new output array,
142	if parameter ``out`` is ``None``.
143	Default: ``"K"``.
144	out (Optional[usm_ndarray]):
145	the array into which the result is written.
146	The data type of `out` must match the expected shape and the
147	expected data type of the result.
148	If ``None`` then a new array is returned. Default: ``None``.
149
150	Returns:
151	usm_ndarray:
152	An array with elements from ``x1`` where ``condition`` is ``True``,
153	and elements from ``x2`` elsewhere.
154
155	The data type of the returned array is determined by applying
156	the Type Promotion Rules to ``x1`` and ``x2``.
157	"""
158	if not isinstance(condition, dpt.usm_ndarray):	1✔
159	raise TypeError(	1✔
160	"Expecting dpctl.tensor.usm_ndarray type, " f"got {type(condition)}"
161	)
162	if order not in ["K", "C", "F", "A"]:	1!
163	order = "K"	×
164	q1, condition_usm_type = condition.sycl_queue, condition.usm_type	1✔
165	q2, x1_usm_type = _get_queue_usm_type(x1)	1✔
166	q3, x2_usm_type = _get_queue_usm_type(x2)	1✔
167	if q2 is None and q3 is None:	1✔
168	exec_q = q1	1✔
169	out_usm_type = condition_usm_type	1✔
170	elif q3 is None:	1✔
171	exec_q = dpctl.utils.get_execution_queue((q1, q2))	1✔
172	if exec_q is None:	1!
173	raise ExecutionPlacementError(	×
174	"Execution placement can not be unambiguously inferred "
175	"from input arguments."
176	)
177	out_usm_type = dpctl.utils.get_coerced_usm_type(	1✔
178	(
179	condition_usm_type,
180	x1_usm_type,
181	)
182	)
183	elif q2 is None:	1✔
184	exec_q = dpctl.utils.get_execution_queue((q1, q3))	1✔
185	if exec_q is None:	1!
186	raise ExecutionPlacementError(	×
187	"Execution placement can not be unambiguously inferred "
188	"from input arguments."
189	)
190	out_usm_type = dpctl.utils.get_coerced_usm_type(	1✔
191	(
192	condition_usm_type,
193	x2_usm_type,
194	)
195	)
196	else:
197	exec_q = dpctl.utils.get_execution_queue((q1, q2, q3))	1✔
198	if exec_q is None:	1✔
199	raise ExecutionPlacementError(	1✔
200	"Execution placement can not be unambiguously inferred "
201	"from input arguments."
202	)
203	out_usm_type = dpctl.utils.get_coerced_usm_type(	1✔
204	(
205	condition_usm_type,
206	x1_usm_type,
207	x2_usm_type,
208	)
209	)
210	dpctl.utils.validate_usm_type(out_usm_type, allow_none=False)	1✔
211	condition_shape = condition.shape	1✔
212	x1_shape = _get_shape(x1)	1✔
213	x2_shape = _get_shape(x2)	1✔
214	if not all(	1!
215	isinstance(s, (tuple, list))
216	for s in (
217	x1_shape,
218	x2_shape,
219	)
220	):
221	raise TypeError(	×
222	"Shape of arguments can not be inferred. "
223	"Arguments are expected to be "
224	"lists, tuples, or both"
225	)
226	try:	1✔
227	res_shape = _broadcast_shape_impl(	1✔
228	[
229	condition_shape,
230	x1_shape,
231	x2_shape,
232	]
233	)
234	except ValueError:	1✔
235	raise ValueError(	1✔
236	"operands could not be broadcast together with shapes "
237	f"{condition_shape}, {x1_shape}, and {x2_shape}"
238	)
239	sycl_dev = exec_q.sycl_device	1✔
240	x1_dtype = _get_dtype(x1, sycl_dev)	1✔
241	x2_dtype = _get_dtype(x2, sycl_dev)	1✔
242	if not all(_validate_dtype(o) for o in (x1_dtype, x2_dtype)):	1✔
243	raise ValueError("Operands have unsupported data types")	1✔
244	x1_dtype, x2_dtype = _resolve_two_weak_types(x1_dtype, x2_dtype, sycl_dev)	1✔
245	out_dtype = _where_result_type(x1_dtype, x2_dtype, sycl_dev)	1✔
246	if out_dtype is None:	1!
247	raise TypeError(	×
248	"function 'where' does not support input "
249	f"types ({x1_dtype}, {x2_dtype}), "
250	"and the inputs could not be safely coerced "
251	"to any supported types according to the casting rule ''safe''."
252	)
253
254	orig_out = out	1✔
255	if out is not None:	1✔
256	if not isinstance(out, dpt.usm_ndarray):	1✔
257	raise TypeError(	1✔
258	"output array must be of usm_ndarray type, got " f"{type(out)}"
259	)
260
261	if not out.flags.writable:	1✔
262	raise ValueError("provided `out` array is read-only")	1✔
263
264	if out.shape != res_shape:	1✔
265	raise ValueError(	1✔
266	"The shape of input and output arrays are "
267	f"inconsistent. Expected output shape is {res_shape}, "
268	f"got {out.shape}"
269	)
270
271	if out_dtype != out.dtype:	1✔
272	raise ValueError(	1✔
273	f"Output array of type {out_dtype} is needed, "
274	f"got {out.dtype}"
275	)
276
277	if dpctl.utils.get_execution_queue((exec_q, out.sycl_queue)) is None:	1✔
278	raise ExecutionPlacementError(	1✔
279	"Input and output allocation queues are not compatible"
280	)
281
282	if ti._array_overlap(condition, out) and not ti._same_logical_tensors(	1✔
283	condition, out
284	):
285	out = dpt.empty_like(out)	1✔
286
287	if isinstance(x1, dpt.usm_ndarray):	1!
288	if (	1✔
289	ti._array_overlap(x1, out)
290	and not ti._same_logical_tensors(x1, out)
291	and x1_dtype == out_dtype
292	):
293	out = dpt.empty_like(out)	1✔
294
295	if isinstance(x2, dpt.usm_ndarray):	1!
296	if (	1✔
297	ti._array_overlap(x2, out)
298	and not ti._same_logical_tensors(x2, out)
299	and x2_dtype == out_dtype
300	):
301	out = dpt.empty_like(out)	1✔
302
303	if order == "A":	1✔
304	order = (	1✔
305	"F"
306	if all(
307	arr.flags.f_contiguous
308	for arr in (
309	condition,
310	x1,
311	x2,
312	)
313	)
314	else "C"
315	)
316	if not isinstance(x1, dpt.usm_ndarray):	1✔
317	x1 = dpt.asarray(x1, dtype=x1_dtype, sycl_queue=exec_q)	1✔
318	if not isinstance(x2, dpt.usm_ndarray):	1✔
319	x2 = dpt.asarray(x2, dtype=x2_dtype, sycl_queue=exec_q)	1✔
320
321	if condition.size == 0:	1✔
322	if out is not None:	1!
323	return out	×
324	else:
325	if order == "K":	1!
326	return _empty_like_triple_orderK(	1✔
327	condition,
328	x1,
329	x2,
330	out_dtype,
331	res_shape,
332	out_usm_type,
333	exec_q,
334	)
335	else:
336	return dpt.empty(	×
337	res_shape,
338	dtype=out_dtype,
339	order=order,
340	usm_type=out_usm_type,
341	sycl_queue=exec_q,
342	)
343
344	_manager = SequentialOrderManager[exec_q]	1✔
345	dep_evs = _manager.submitted_events	1✔
346	if x1_dtype != out_dtype:	1✔
347	if order == "K":	1✔
348	_x1 = _empty_like_orderK(x1, out_dtype)	1✔
349	else:
350	_x1 = dpt.empty_like(x1, dtype=out_dtype, order=order)	1✔
351	ht_copy1_ev, copy1_ev = ti._copy_usm_ndarray_into_usm_ndarray(	1✔
352	src=x1, dst=_x1, sycl_queue=exec_q, depends=dep_evs
353	)
354	x1 = _x1	1✔
355	_manager.add_event_pair(ht_copy1_ev, copy1_ev)	1✔
356
357	if x2_dtype != out_dtype:	1✔
358	if order == "K":	1✔
359	_x2 = _empty_like_orderK(x2, out_dtype)	1✔
360	else:
361	_x2 = dpt.empty_like(x2, dtype=out_dtype, order=order)	1✔
362	ht_copy2_ev, copy2_ev = ti._copy_usm_ndarray_into_usm_ndarray(	1✔
363	src=x2, dst=_x2, sycl_queue=exec_q, depends=dep_evs
364	)
365	x2 = _x2	1✔
366	_manager.add_event_pair(ht_copy2_ev, copy2_ev)	1✔
367
368	if out is None:	1✔
369	if order == "K":	1✔
370	out = _empty_like_triple_orderK(	1✔
371	condition, x1, x2, out_dtype, res_shape, out_usm_type, exec_q
372	)
373	else:
374	out = dpt.empty(	1✔
375	res_shape,
376	dtype=out_dtype,
377	order=order,
378	usm_type=out_usm_type,
379	sycl_queue=exec_q,
380	)
381
382	if condition_shape != res_shape:	1✔
383	condition = dpt.broadcast_to(condition, res_shape)	1✔
384	if x1_shape != res_shape:	1✔
385	x1 = dpt.broadcast_to(x1, res_shape)	1✔
386	if x2_shape != res_shape:	1✔
387	x2 = dpt.broadcast_to(x2, res_shape)	1✔
388
389	dep_evs = _manager.submitted_events	1✔
390	hev, where_ev = ti._where(	1✔
391	condition=condition,
392	x1=x1,
393	x2=x2,
394	dst=out,
395	sycl_queue=exec_q,
396	depends=dep_evs,
397	)
398	_manager.add_event_pair(hev, where_ev)	1✔
399	if not (orig_out is None or orig_out is out):	1✔
400	# Copy the out data from temporary buffer to original memory
401	ht_copy_out_ev, cpy_ev = ti._copy_usm_ndarray_into_usm_ndarray(	1✔
402	src=out,
403	dst=orig_out,
404	sycl_queue=exec_q,
405	depends=[where_ev],
406	)
407	_manager.add_event_pair(ht_copy_out_ev, cpy_ev)	1✔
408	out = orig_out	1✔
409
410	return out	1✔

IntelPython / dpctl / 11468253350

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous