Allow x to be a scalar in isin and remove assume_unique

Author: ndgrigorian

dpctl/tensor/_set_functions.py

@@ -21,7 +21,12 @@
 import dpctl.utils as du
 
 from ._copy_utils import _empty_like_orderK
-from ._scalar_utils import _get_dtype, _get_queue_usm_type, _validate_dtype
+from ._scalar_utils import (
+    _get_dtype,
+    _get_queue_usm_type,
+    _get_shape,
+    _validate_dtype,
+)
 from ._tensor_elementwise_impl import _not_equal, _subtract
 from ._tensor_impl import (
     _copy_usm_ndarray_into_usm_ndarray,
@@ -38,7 +43,10 @@
     _searchsorted_left,
     _sort_ascending,
 )
-from ._type_utils import _resolve_weak_types_all_py_ints
+from ._type_utils import (
+    _resolve_weak_types_all_py_ints,
+    _to_device_supported_dtype,
+)
 
 __all__ = [
     "isin",
@@ -632,21 +640,17 @@ def unique_all(x: dpt.usm_ndarray) -> UniqueAllResult:
     )
 
 
-def isin(x, test_elements, /, *, assume_unique=False, invert=False):
+def isin(x, test_elements, /, *, invert=False):
     """
     Tests `x in test_elements` for each element of `x`. Returns a boolean array
     with the same shape as `x` that is `True` where the element is in
     `test_elements`, `False` otherwise.
 
     Args:
-        x (usm_ndarray):
-            input array.
+        x (Union[usm_ndarray, bool, int, float, complex]):
+            input element or elements.
         test_elements (Union[usm_ndarray, bool, int, float, complex]):
             elements against which to test each value of `x`.
-        assume_unique (Optional[bool]):
-            if `True`, the input arrays are both assumed to be unique, which
-            currently has no effect.
-            Default: `False`.
         invert (Optional[bool]):
             if `True`, the output results are inverted, i.e., are equivalent to
             testing `x not in test_elements` for each element of `x`.
@@ -657,11 +661,19 @@ def isin(x, test_elements, /, *, assume_unique=False, invert=False):
             an array of the inclusion test results. The returned array has a
             boolean data type and the same shape as `x`.
     """
-    if not isinstance(x, dpt.usm_ndarray):
-        raise TypeError(f"Expected dpctl.tensor.usm_ndarray, got {type(x)}")
-    q1, x_usm_type = x.sycl_queue, x.usm_type
+    q1, x_usm_type = _get_queue_usm_type(x)
     q2, test_usm_type = _get_queue_usm_type(test_elements)
-    if q2 is None:
+    if q1 is None and q2 is None:
+        raise du.ExecutionPlacementError(
+            "Execution placement can not be unambiguously inferred "
+            "from input arguments. "
+            "One of the arguments must represent USM allocation and "
+            "expose `__sycl_usm_array_interface__` property"
+        )
+    if q1 is None:
+        exec_q = q2
+        res_usm_type = test_usm_type
+    elif q2 is None:
         exec_q = q1
         res_usm_type = x_usm_type
     else:
@@ -680,45 +692,60 @@ def isin(x, test_elements, /, *, assume_unique=False, invert=False):
     dpctl.utils.validate_usm_type(res_usm_type, allow_none=False)
     sycl_dev = exec_q.sycl_device
 
+    x_dt = _get_dtype(x, sycl_dev)
+    test_dt = _get_dtype(test_elements, sycl_dev)
+    if not all(_validate_dtype(dt) for dt in (x_dt, test_dt)):
+        raise ValueError("Operands have unsupported data types")
+
+    x_sh = _get_shape(x)
     if isinstance(test_elements, dpt.usm_ndarray) and test_elements.size == 0:
         if invert:
-            return dpt.ones_like(x, dtype=dpt.bool, usm_type=res_usm_type)
+            return dpt.ones(
+                x_sh, dtype=dpt.bool, usm_type=res_usm_type, sycl_queue=exec_q
+            )
         else:
-            return dpt.zeros_like(x, dtype=dpt.bool, usm_type=res_usm_type)
+            return dpt.zeros(
+                x_sh, dtype=dpt.bool, usm_type=res_usm_type, sycl_queue=exec_q
+            )
 
-    x_dt = x.dtype
-    test_dt = _get_dtype(test_elements, sycl_dev)
-    if not _validate_dtype(test_dt):
-        raise ValueError("`test_elements` has unsupported dtype")
+    dt1, dt2 = _resolve_weak_types_all_py_ints(x_dt, test_dt, sycl_dev)
+    dt = _to_device_supported_dtype(dpt.result_type(dt1, dt2), sycl_dev)
+
+    if not isinstance(x, dpt.usm_ndarray):
+        x_arr = dpt.asarray(
+            x, dtype=dt1, usm_type=res_usm_type, sycl_queue=exec_q
+        )
+    else:
+        x_arr = x
+
+    if not isinstance(test_elements, dpt.usm_ndarray):
+        test_arr = dpt.asarray(
+            test_elements, dtype=dt2, usm_type=res_usm_type, sycl_queue=exec_q
+        )
+    else:
+        test_arr = test_elements
 
     _manager = du.SequentialOrderManager[exec_q]
     dep_evs = _manager.submitted_events
 
-    dt1, dt2 = _resolve_weak_types_all_py_ints(x_dt, test_dt, sycl_dev)
-    dt = dpt.result_type(dt1, dt2)
-
     if x_dt != dt:
-        x_buf = _empty_like_orderK(x, dt)
+        x_buf = _empty_like_orderK(x_arr, dt, res_usm_type, sycl_dev)
         ht_ev, ev = _copy_usm_ndarray_into_usm_ndarray(
-            src=x, dst=x_buf, sycl_queue=exec_q, depends=dep_evs
+            src=x_arr, dst=x_buf, sycl_queue=exec_q, depends=dep_evs
         )
         _manager.add_event_pair(ht_ev, ev)
     else:
-        x_buf = x
+        x_buf = x_arr
 
-    if not isinstance(test_elements, dpt.usm_ndarray):
-        test_buf = dpt.asarray(
-            test_elements, dtype=dt, usm_type=res_usm_type, sycl_queue=exec_q
-        )
-    elif test_dt != dt:
+    if test_dt != dt:
         # copy into C-contiguous memory, because the array will be flattened
-        test_buf = dpt.empty_like(test_elements, dtype=dt, order="C")
+        test_buf = dpt.empty_like(test_arr, dtype=dt, order="C")
         ht_ev, ev = _copy_usm_ndarray_into_usm_ndarray(
-            src=test_elements, dst=test_buf, sycl_queue=exec_q, depends=dep_evs
+            src=test_arr, dst=test_buf, sycl_queue=exec_q, depends=dep_evs
         )
         _manager.add_event_pair(ht_ev, ev)
     else:
-        test_buf = test_elements
+        test_buf = test_arr
 
     test_buf = dpt.reshape(test_buf, -1)
     test_buf = dpt.sort(test_buf)