MAINT: searchsorted: move delegation to _delegation.py

Author: mdhaber

src/array_api_extra/__init__.py

@@ -12,6 +12,7 @@
     one_hot,
     pad,
     partition,
+    searchsorted,
     setdiff1d,
     sinc,
     union1d,
@@ -23,7 +24,6 @@
     default_dtype,
     kron,
     nunique,
-    searchsorted,
 )
 from ._lib._lazy import lazy_apply
 

src/array_api_extra/_delegation.py

@@ -27,6 +27,7 @@
     "nan_to_num",
     "one_hot",
     "pad",
+    "searchsorted",
     "sinc",
 ]
 
@@ -632,6 +633,76 @@ def pad(
     return _funcs.pad(x, pad_width, constant_values=constant_values, xp=xp)
 
 
+def searchsorted(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    side: Literal["left", "right"] = "left",
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Find indices where elements should be inserted to maintain order.
+
+    Find the indices into a sorted array ``x1`` such that if the elements in ``x2``
+    were inserted before the indices, the resulting array would remain sorted.
+
+    Parameters
+    ----------
+    x1 : Array
+        Input array. Should have a real-valued data type. Must be sorted in ascending
+        order along the last axis.
+    x2 : Array
+        Array containing search values. Should have a real-valued data type. Must have
+        the same shape as ``x1`` except along the last axis.
+    side : {'left', 'right'}, optional
+        Argument controlling which index is returned if an element of ``x2`` is equal to
+        one or more elements of ``x1``: ``'left'`` returns the index of the first of
+        these elements; ``'right'`` returns the next index after the last of these
+        elements. Default: ``'left'``.
+    xp : array_namespace, optional
+        The standard-compatible namespace for the array arguments. Default: infer.
+
+    Returns
+    -------
+    Array: integer array
+        An array of indices with the same shape as ``x2``.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> x = xp.asarray([11, 12, 13, 13, 14, 15])
+    >>> xpx.searchsorted(x, xp.asarray([10, 11.5, 14.5, 16]), xp=xp)
+    Array([0, 1, 5, 6], dtype=array_api_strict.int64)
+    >>> xpx.searchsorted(x, xp.asarray(13), xp=xp)
+    Array(2, dtype=array_api_strict.int64)
+    >>> xpx.searchsorted(x, xp.asarray(13), side='right', xp=xp)
+    Array(4, dtype=array_api_strict.int64)
+
+    `searchsorted` is vectorized along the last axis.
+
+    >>> x1 = xp.asarray([[1., 2., 3., 4.], [5., 6., 7., 8.]])
+    >>> x2 = xp.asarray([[1.1, 3.3], [6.6, 8.8]])
+    >>> xpx.searchsorted(x1, x2, xp=xp)
+    Array([[1, 3],
+           [2, 4]], dtype=array_api_strict.int64)
+    """
+    if xp is None:
+        xp = array_namespace(x1, x2)
+
+    xp_default_int = _funcs.default_dtype(xp, kind="integral")
+    y_0d = xp.asarray(x2).ndim == 0
+    x_1d = x1.ndim <= 1
+
+    if x_1d or is_torch_namespace(xp):
+        x2 = xp.reshape(x2, ()) if (y_0d and x_1d) else x2
+        out = xp.searchsorted(x1, x2, side=side)
+        return xp.astype(out, xp_default_int, copy=False)
+
+    return _funcs.searchsorted(x1, x2, side=side, xp=xp)
+
+
 def setdiff1d(
     x1: Array | complex,
     x2: Array | complex,

src/array_api_extra/_lib/_funcs.py

@@ -12,7 +12,6 @@
     array_namespace,
     is_dask_namespace,
     is_jax_array,
-    is_torch_namespace,
 )
 from ._utils._helpers import (
     asarrays,
@@ -677,65 +676,10 @@ def searchsorted(
     /,
     *,
     side: Literal["left", "right"] = "left",
-    xp: ModuleType | None = None,
+    xp: ModuleType,
 ) -> Array:
-    """
-    Find indices where elements should be inserted to maintain order.
-
-    Find the indices into a sorted array ``x1`` such that if the elements in ``x2``
-    were inserted before the indices, the resulting array would remain sorted.
-
-    Parameters
-    ----------
-    x1 : Array
-        Input array. Should have a real-valued data type. Must be sorted in ascending
-        order along the last axis.
-    x2 : Array
-        Array containing search values. Should have a real-valued data type. Must have
-        the same shape as ``x1`` except along the last axis.
-    side : {'left', 'right'}, optional
-        Argument controlling which index is returned if an element of ``x2`` is equal to
-        one or more elements of ``x1``: ``'left'`` returns the index of the first of
-        these elements; ``'right'`` returns the next index after the last of these
-        elements. Default: ``'left'``.
-    xp : array_namespace, optional
-        The standard-compatible namespace for the array arguments. Default: infer.
-
-    Returns
-    -------
-    Array: integer array
-        An array of indices with the same shape as ``x2``.
-
-    Examples
-    --------
-    >>> import array_api_strict as xp
-    >>> import array_api_extra as xpx
-    >>> x = xp.asarray([11, 12, 13, 13, 14, 15])
-    >>> xpx.searchsorted(x, xp.asarray([10, 11.5, 14.5, 16]), xp=xp)
-    Array([0, 1, 5, 6], dtype=array_api_strict.int64)
-    >>> xpx.searchsorted(x, xp.asarray(13), xp=xp)
-    Array(2, dtype=array_api_strict.int64)
-    >>> xpx.searchsorted(x, xp.asarray(13), side='right', xp=xp)
-    Array(4, dtype=array_api_strict.int64)
-
-    `searchsorted` is vectorized along the last axis.
-
-    >>> x1 = xp.asarray([[1., 2., 3., 4.], [5., 6., 7., 8.]])
-    >>> x2 = xp.asarray([[1.1, 3.3], [6.6, 8.8]])
-    >>> xpx.searchsorted(x1, x2, xp=xp)
-    Array([[1, 3],
-           [2, 4]], dtype=array_api_strict.int64)
-    """
-    xp = array_namespace(x1, x2) if xp is None else xp
-    xp_default_int = xp.asarray(1).dtype
-    y_0d = xp.asarray(x2).ndim == 0
-    x_1d = x1.ndim <= 1
-
-    if x_1d or is_torch_namespace(xp):
-        x2 = xp.reshape(x2, ()) if (y_0d and x_1d) else x2
-        out = xp.searchsorted(x1, x2, side=side)
-        return xp.astype(out, xp_default_int, copy=False)
-
+    # numpydoc ignore=PR01,RT01
+    """See docstring in `array_api_extra._delegation.py`."""
     a = xp.full(x2.shape, 0, device=_compat.device(x1))
 
     if x1.shape[-1] == 0:
@@ -757,7 +701,7 @@ def searchsorted(
 
     out = xp.where(compare(x2, xp.min(x1, axis=-1, keepdims=True)), 0, b)
     out = xp.where(xp.isnan(x2), x1.shape[-1], out) if side == "right" else out
-    return xp.astype(out, xp_default_int, copy=False)
+    return xp.astype(out, default_dtype(xp, kind="integral"), copy=False)
 
 
 def setdiff1d(