Merge pull request #2098 from IntelPython/feature/searchsorted-based-isin

Implement `tensor.isin`

Author: ndgrigorian

docs/doc_sources/api_reference/dpctl/tensor.set_functions.rst

@@ -8,6 +8,7 @@ Set Functions
 .. autosummary::
     :toctree: generated
 
+    isin
     unique_all
     unique_counts
     unique_inverse

dpctl/tensor/CMakeLists.txt

@@ -112,6 +112,7 @@ set(_reduction_sources
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/reductions/sum.cpp
 )
 set(_sorting_sources
+    ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/isin.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/merge_sort.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/merge_argsort.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/libtensor/source/sorting/radix_sort.cpp

dpctl/tensor/__init__.py

@@ -198,6 +198,7 @@
 )
 from ._searchsorted import searchsorted
 from ._set_functions import (
+    isin,
     unique_all,
     unique_counts,
     unique_inverse,
@@ -394,4 +395,5 @@
     "top_k",
     "dldevice_to_sycl_device",
     "sycl_device_to_dldevice",
+    "isin",
 ]

dpctl/tensor/_clip.py

@@ -23,16 +23,16 @@
     _empty_like_pair_orderK,
     _empty_like_triple_orderK,
 )
-from dpctl.tensor._elementwise_common import (
+from dpctl.tensor._manipulation_functions import _broadcast_shape_impl
+from dpctl.tensor._type_utils import _can_cast
+from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
+
+from ._scalar_utils import (
     _get_dtype,
     _get_queue_usm_type,
     _get_shape,
     _validate_dtype,
 )
-from dpctl.tensor._manipulation_functions import _broadcast_shape_impl
-from dpctl.tensor._type_utils import _can_cast
-from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
-
 from ._type_utils import (
     _resolve_one_strong_one_weak_types,
     _resolve_one_strong_two_weak_types,

dpctl/tensor/_elementwise_common.py

@@ -14,32 +14,27 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import numbers
-
-import numpy as np
-
 import dpctl
-import dpctl.memory as dpm
 import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
 from dpctl.tensor._manipulation_functions import _broadcast_shape_impl
-from dpctl.tensor._usmarray import _is_object_with_buffer_protocol as _is_buffer
 from dpctl.utils import ExecutionPlacementError, SequentialOrderManager
 
 from ._copy_utils import _empty_like_orderK, _empty_like_pair_orderK
+from ._scalar_utils import (
+    _get_dtype,
+    _get_queue_usm_type,
+    _get_shape,
+    _validate_dtype,
+)
 from ._type_utils import (
-    WeakBooleanType,
-    WeakComplexType,
-    WeakFloatingType,
-    WeakIntegralType,
     _acceptance_fn_default_binary,
     _acceptance_fn_default_unary,
     _all_data_types,
     _find_buf_dtype,
     _find_buf_dtype2,
     _find_buf_dtype_in_place_op,
     _resolve_weak_types,
-    _to_device_supported_dtype,
 )
 
 
@@ -289,78 +284,6 @@ def __call__(self, x, /, *, out=None, order="K"):
         return out
 
 
-def _get_queue_usm_type(o):
-    """Return SYCL device where object `o` allocated memory, or None."""
-    if isinstance(o, dpt.usm_ndarray):
-        return o.sycl_queue, o.usm_type
-    elif hasattr(o, "__sycl_usm_array_interface__"):
-        try:
-            m = dpm.as_usm_memory(o)
-            return m.sycl_queue, m.get_usm_type()
-        except Exception:
-            return None, None
-    return None, None
-
-
-def _get_dtype(o, dev):
-    if isinstance(o, dpt.usm_ndarray):
-        return o.dtype
-    if hasattr(o, "__sycl_usm_array_interface__"):
-        return dpt.asarray(o).dtype
-    if _is_buffer(o):
-        host_dt = np.array(o).dtype
-        dev_dt = _to_device_supported_dtype(host_dt, dev)
-        return dev_dt
-    if hasattr(o, "dtype"):
-        dev_dt = _to_device_supported_dtype(o.dtype, dev)
-        return dev_dt
-    if isinstance(o, bool):
-        return WeakBooleanType(o)
-    if isinstance(o, int):
-        return WeakIntegralType(o)
-    if isinstance(o, float):
-        return WeakFloatingType(o)
-    if isinstance(o, complex):
-        return WeakComplexType(o)
-    return np.object_
-
-
-def _validate_dtype(dt) -> bool:
-    return isinstance(
-        dt,
-        (WeakBooleanType, WeakIntegralType, WeakFloatingType, WeakComplexType),
-    ) or (
-        isinstance(dt, dpt.dtype)
-        and dt
-        in [
-            dpt.bool,
-            dpt.int8,
-            dpt.uint8,
-            dpt.int16,
-            dpt.uint16,
-            dpt.int32,
-            dpt.uint32,
-            dpt.int64,
-            dpt.uint64,
-            dpt.float16,
-            dpt.float32,
-            dpt.float64,
-            dpt.complex64,
-            dpt.complex128,
-        ]
-    )
-
-
-def _get_shape(o):
-    if isinstance(o, dpt.usm_ndarray):
-        return o.shape
-    if _is_buffer(o):
-        return memoryview(o).shape
-    if isinstance(o, numbers.Number):
-        return tuple()
-    return getattr(o, "shape", tuple())
-
-
 class BinaryElementwiseFunc:
     """
     Class that implements binary element-wise functions.

dpctl/tensor/_scalar_utils.py

@@ -0,0 +1,111 @@
+#                       Data Parallel Control (dpctl)
+#
+#  Copyright 2020-2025 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 numbers
+
+import numpy as np
+
+import dpctl.memory as dpm
+import dpctl.tensor as dpt
+from dpctl.tensor._usmarray import _is_object_with_buffer_protocol as _is_buffer
+
+from ._type_utils import (
+    WeakBooleanType,
+    WeakComplexType,
+    WeakFloatingType,
+    WeakIntegralType,
+    _to_device_supported_dtype,
+)
+
+
+def _get_queue_usm_type(o):
+    """Return SYCL device where object `o` allocated memory, or None."""
+    if isinstance(o, dpt.usm_ndarray):
+        return o.sycl_queue, o.usm_type
+    elif hasattr(o, "__sycl_usm_array_interface__"):
+        try:
+            m = dpm.as_usm_memory(o)
+            return m.sycl_queue, m.get_usm_type()
+        except Exception:
+            return None, None
+    return None, None
+
+
+def _get_dtype(o, dev):
+    if isinstance(o, dpt.usm_ndarray):
+        return o.dtype
+    if hasattr(o, "__sycl_usm_array_interface__"):
+        return dpt.asarray(o).dtype
+    if _is_buffer(o):
+        host_dt = np.array(o).dtype
+        dev_dt = _to_device_supported_dtype(host_dt, dev)
+        return dev_dt
+    if hasattr(o, "dtype"):
+        dev_dt = _to_device_supported_dtype(o.dtype, dev)
+        return dev_dt
+    if isinstance(o, bool):
+        return WeakBooleanType(o)
+    if isinstance(o, int):
+        return WeakIntegralType(o)
+    if isinstance(o, float):
+        return WeakFloatingType(o)
+    if isinstance(o, complex):
+        return WeakComplexType(o)
+    return np.object_
+
+
+def _validate_dtype(dt) -> bool:
+    return isinstance(
+        dt,
+        (WeakBooleanType, WeakIntegralType, WeakFloatingType, WeakComplexType),
+    ) or (
+        isinstance(dt, dpt.dtype)
+        and dt
+        in [
+            dpt.bool,
+            dpt.int8,
+            dpt.uint8,
+            dpt.int16,
+            dpt.uint16,
+            dpt.int32,
+            dpt.uint32,
+            dpt.int64,
+            dpt.uint64,
+            dpt.float16,
+            dpt.float32,
+            dpt.float64,
+            dpt.complex64,
+            dpt.complex128,
+        ]
+    )
+
+
+def _get_shape(o):
+    if isinstance(o, dpt.usm_ndarray):
+        return o.shape
+    if _is_buffer(o):
+        return memoryview(o).shape
+    if isinstance(o, numbers.Number):
+        return tuple()
+    return getattr(o, "shape", tuple())
+
+
+__all__ = [
+    "_get_dtype",
+    "_get_queue_usm_type",
+    "_get_shape",
+    "_validate_dtype",
+]

dpctl/tensor/_search_functions.py

@@ -17,16 +17,16 @@
 import dpctl
 import dpctl.tensor as dpt
 import dpctl.tensor._tensor_impl as ti
-from dpctl.tensor._elementwise_common import (
+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 ._scalar_utils 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,