permit Python scalars as second argument to dpt.searchsorted

Author: ndgrigorian

dpctl/tensor/_searchsorted.py

@@ -1,23 +1,29 @@
 from typing import Literal, Union
 
 import dpctl
+import dpctl.tensor as dpt
 import dpctl.utils as du
 
 from ._copy_utils import _empty_like_orderK
 from ._ctors import empty
+from ._scalar_utils import _get_dtype, _get_queue_usm_type, _validate_dtype
 from ._tensor_impl import _copy_usm_ndarray_into_usm_ndarray as ti_copy
 from ._tensor_impl import _take as ti_take
 from ._tensor_impl import (
     default_device_index_type as ti_default_device_index_type,
 )
 from ._tensor_sorting_impl import _searchsorted_left, _searchsorted_right
-from ._type_utils import isdtype, result_type
+from ._type_utils import (
+    _resolve_weak_types_all_py_ints,
+    _to_device_supported_dtype,
+    isdtype,
+)
 from ._usmarray import usm_ndarray
 
 
 def searchsorted(
     x1: usm_ndarray,
-    x2: usm_ndarray,
+    x2: Union[usm_ndarray, int, float, complex, bool],
     /,
     *,
     side: Literal["left", "right"] = "left",
@@ -34,8 +40,8 @@ def searchsorted(
             input array. Must be a one-dimensional array. If `sorter` is
             `None`, must be sorted in ascending order; otherwise, `sorter` must
             be an array of indices that sort `x1` in ascending order.
-        x2 (usm_ndarray):
-            array containing search values.
+        x2 (Union[usm_ndarray, bool, int, float, complex]):
+            search value or values.
         side (Literal["left", "right]):
             argument controlling which index is returned if a value lands
             exactly on an edge. If `x2` is an array of rank `N` where
@@ -56,8 +62,6 @@ def searchsorted(
     """
     if not isinstance(x1, usm_ndarray):
         raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x1)}")
-    if not isinstance(x2, usm_ndarray):
-        raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x2)}")
     if sorter is not None and not isinstance(sorter, usm_ndarray):
         raise TypeError(
             f"Expected dpctl.tensor.usm_ndarray, got {type(sorter)}"
@@ -69,23 +73,39 @@ def searchsorted(
             "Expected either 'left' or 'right'"
         )
 
-    if sorter is None:
-        q = du.get_execution_queue([x1.sycl_queue, x2.sycl_queue])
-    else:
-        q = du.get_execution_queue(
-            [x1.sycl_queue, x2.sycl_queue, sorter.sycl_queue]
-        )
+    q1, x1_usm_type = x1.sycl_queue, x1.usm_type
+    q2, x2_usm_type = _get_queue_usm_type(x2)
+    q3 = sorter.sycl_queue if sorter is not None else None
+    q = du.get_execution_queue(tuple(q for q in (q1, q2, q3) if q is not None))
     if q is None:
         raise du.ExecutionPlacementError(
             "Execution placement can not be unambiguously "
             "inferred from input arguments."
         )
 
+    res_usm_type = du.get_coerced_usm_type(
+        tuple(
+            ut
+            for ut in (
+                x1_usm_type,
+                x2_usm_type,
+            )
+            if ut is not None
+        )
+    )
+    du.validate_usm_type(res_usm_type, allow_none=False)
+    sycl_dev = q.sycl_device
+
     if x1.ndim != 1:
         raise ValueError("First argument array must be one-dimensional")
 
     x1_dt = x1.dtype
-    x2_dt = x2.dtype
+    x2_dt = _get_dtype(x2, sycl_dev)
+    if not _validate_dtype(x2_dt):
+        raise ValueError(
+            "dpt.searchsorted search value argument has "
+            f"unsupported data type {x2_dt}"
+        )
 
     _manager = du.SequentialOrderManager[q]
     dep_evs = _manager.submitted_events
@@ -100,7 +120,7 @@ def searchsorted(
                 "Sorter array must be one-dimension with the same "
                 "shape as the first argument array"
             )
-        res = empty(x1.shape, dtype=x1_dt, usm_type=x1.usm_type, sycl_queue=q)
+        res = empty(x1.shape, dtype=x1_dt, usm_type=x1_usm_type, sycl_queue=q)
         ind = (sorter,)
         axis = 0
         wrap_out_of_bound_indices_mode = 0
@@ -116,29 +136,28 @@ def searchsorted(
         x1 = res
         _manager.add_event_pair(ht_ev, ev)
 
-    if x1_dt != x2_dt:
-        dt = result_type(x1, x2)
-        if x1_dt != dt:
-            x1_buf = _empty_like_orderK(x1, dt)
-            dep_evs = _manager.submitted_events
-            ht_ev, ev = ti_copy(
-                src=x1, dst=x1_buf, sycl_queue=q, depends=dep_evs
-            )
-            _manager.add_event_pair(ht_ev, ev)
-            x1 = x1_buf
-        if x2_dt != dt:
-            x2_buf = _empty_like_orderK(x2, dt)
-            dep_evs = _manager.submitted_events
-            ht_ev, ev = ti_copy(
-                src=x2, dst=x2_buf, sycl_queue=q, depends=dep_evs
-            )
-            _manager.add_event_pair(ht_ev, ev)
-            x2 = x2_buf
+    dt1, dt2 = _resolve_weak_types_all_py_ints(x1_dt, x2_dt, sycl_dev)
+    dt = _to_device_supported_dtype(dpt.result_type(dt1, dt2), sycl_dev)
+
+    if x1_dt != dt:
+        x1_buf = _empty_like_orderK(x1, dt)
+        dep_evs = _manager.submitted_events
+        ht_ev, ev = ti_copy(src=x1, dst=x1_buf, sycl_queue=q, depends=dep_evs)
+        _manager.add_event_pair(ht_ev, ev)
+        x1 = x1_buf
+
+    if not isinstance(x2, usm_ndarray):
+        x2 = dpt.asarray(x2, dtype=dt2, usm_type=res_usm_type, sycl_queue=q)
+    if x2.dtype != dt:
+        x2_buf = _empty_like_orderK(x2, dt)
+        dep_evs = _manager.submitted_events
+        ht_ev, ev = ti_copy(src=x2, dst=x2_buf, sycl_queue=q, depends=dep_evs)
+        _manager.add_event_pair(ht_ev, ev)
+        x2 = x2_buf
 
-    dst_usm_type = du.get_coerced_usm_type([x1.usm_type, x2.usm_type])
     index_dt = ti_default_device_index_type(q)
 
-    dst = _empty_like_orderK(x2, index_dt, usm_type=dst_usm_type)
+    dst = _empty_like_orderK(x2, index_dt, usm_type=res_usm_type)
 
     dep_evs = _manager.submitted_events
     if side == "left":