ENH: Add vectorized searchsorted (#531)

Author: mdhaber

src/array_api_extra/__init__.py

@@ -12,6 +12,7 @@
     one_hot,
     pad,
     partition,
+    searchsorted,
     setdiff1d,
     sinc,
     union1d,
@@ -49,6 +50,7 @@
     "one_hot",
     "pad",
     "partition",
+    "searchsorted",
     "setdiff1d",
     "sinc",
     "union1d",

src/array_api_extra/_delegation.py

@@ -27,6 +27,7 @@
     "nan_to_num",
     "one_hot",
     "pad",
+    "searchsorted",
     "sinc",
 ]
 
@@ -632,6 +633,85 @@ 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.
+
+    The behavior of this function is similar to that of `array_api.searchsorted`,
+    but it relaxes the requirement that `x1` must be one-dimensional.
+    This function is vectorized, treating slices along the last axis
+    as elements and preceding axes as batch (or "loop") dimensions.
+
+    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)
+
+    if side not in {"left", "right"}:
+        message = "`side` must be either 'left' or 'right'."
+        raise ValueError(message)
+
+    xp_default_int = _funcs.default_dtype(xp, kind="integral")
+    x2_0d = x2.ndim == 0
+    x1_1d = x1.ndim <= 1
+
+    if x1_1d or is_torch_namespace(xp):
+        x2 = xp.reshape(x2, ()) if (x2_0d and x1_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

@@ -8,7 +8,11 @@
 
 from ._at import at
 from ._utils import _compat, _helpers
-from ._utils._compat import array_namespace, is_dask_namespace, is_jax_array
+from ._utils._compat import (
+    array_namespace,
+    is_dask_namespace,
+    is_jax_array,
+)
 from ._utils._helpers import (
     asarrays,
     capabilities,
@@ -28,6 +32,7 @@
     "kron",
     "nunique",
     "pad",
+    "searchsorted",
     "setdiff1d",
     "sinc",
 ]
@@ -693,6 +698,40 @@ def pad(
     return at(padded, tuple(slices)).set(x)
 
 
+def searchsorted(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    side: Literal["left", "right"] = "left",
+    xp: ModuleType,
+) -> Array:
+    # 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:
+        return a
+
+    n = xp.count_nonzero(~xp.isnan(x1), axis=-1, keepdims=True)
+    b = xp.broadcast_to(n, x2.shape)
+
+    compare = xp.less_equal if side == "left" else xp.less
+
+    # while xp.any(b - a > 1):
+    # refactored to for loop with ~log2(n) iterations for JAX JIT
+    for _ in range(int(math.log2(x1.shape[-1])) + 1):  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
+        c = (a + b) // 2
+        x0 = xp.take_along_axis(x1, c, axis=-1)
+        j = compare(x2, x0)
+        b = xp.where(j, c, b)
+        a = xp.where(j, a, c)
+
+    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, default_dtype(xp, kind="integral"), copy=False)
+
+
 def setdiff1d(
     x1: Array | complex,
     x2: Array | complex,