Add searchsorted

As far as I can tell, except for the dtype restriction, the standard is the
same as NumPy.

Author: asmeurer

array_api_strict/__init__.py

@@ -286,9 +286,9 @@
 
 __all__ += ["concat", "expand_dims", "flip", "moveaxis", "permute_dims", "repeat", "reshape", "roll", "squeeze", "stack"]
 
-from ._searching_functions import argmax, argmin, nonzero, where
+from ._searching_functions import argmax, argmin, nonzero, searchsorted, where
 
-__all__ += ["argmax", "argmin", "nonzero", "where"]
+__all__ += ["argmax", "argmin", "nonzero", "searchsorted", "where"]
 
 from ._set_functions import unique_all, unique_counts, unique_inverse, unique_values
 

array_api_strict/_searching_functions.py

@@ -2,11 +2,11 @@
 
 from ._array_object import Array
 from ._dtypes import _result_type, _real_numeric_dtypes
-from ._flags import requires_data_dependent_shapes
+from ._flags import requires_data_dependent_shapes, requires_api_version
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from typing import Optional, Tuple
+    from typing import Literal, Optional, Tuple
 
 import numpy as np
 
@@ -45,6 +45,26 @@ def nonzero(x: Array, /) -> Tuple[Array, ...]:
         raise ValueError("nonzero is not allowed on 0-dimensional arrays")
     return tuple(Array._new(i) for i in np.nonzero(x._array))
 
+@requires_api_version('2023.12')
+def searchsorted(
+    x1: Array,
+    x2: Array,
+    /,
+    *,
+    side: Literal["left", "right"] = "left",
+    sorter: Optional[Array] = None,
+) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.searchsorted <numpy.searchsorted>`.
+
+    See its docstring for more information.
+    """
+    if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in searchsorted")
+    sorter = sorter._array if sorter is not None else None
+    # TODO: The sort order of nans and signed zeros is implementation
+    # dependent. Should we error/warn if they are present?
+    return Array._new(np.searchsorted(x1._array, x2._array, side=side, sorter=sorter))
 
 def where(condition: Array, x1: Array, x2: Array, /) -> Array:
     """