Repository: google/jax

Issue:
`jnp.searchsorted` returns different values to `np.searchsorted`
```python
import numpy as np
import jax.numpy as jnp

x = np.arange(10, dtype=float)
x[4:] = jnp.nan
np.searchsorted(x, 0)  # 0
jnp.searchsorted(x, 0)  # DeviceArray(10, dtype=int32)
```
It doesn't look like `jnp.searchsorted` handles NaNs correctly.


Discussion:
I think this is expected: for `searchsorted` to work correctly, the array must be sorted in increasing order. But for your input I get this:
```python
>>> x[:-1] < x[1:]
array([ True,  True,  True, False, False, False, False, False, False])
```
which implies the array is not in a sorted order, which means the output is ill-defined. We have not made any effort to ensure that JAX's choices for ill-defined outputs match numpy's, as the details depend on the implementation, and so it would be quite difficult to do so (side note: even for `sort()`, JAX treats NaNs differently than NumPy does; see #6443)

What do you think? Is this something you can work around?
Thinking about it more... we could consider this defined behavior in that `sort()` is defined for `NaN` (and `inf`) entries. With that in mind, perhaps we should require that the index returned by `searchsorted` is that that would leave the array unchanged under JAX's definition of `sort()`. This could be done fairly easily, by first calling `_float_to_int_for_sort` on the entries, as is done when `sort()` is called on a float array: https://github.com/google/jax/blob/f96761cdc861ed6ecaa3bd1b84ca88b257d1b601/jax/_src/lax/lax.py#L3649

That still would not exactly match NumPy's behavior because of the handling of negative NaNs (cf. #6443) but it would definitely be more sensible than the current behavior. What do you think?
(side-note: this is quite timely becaue I've been working on `searchsorted` all this afternoon: #9108)
I think your proposal to handle NaNs as per `jnp.sort` sounds like a good one. Thanks!

This is the choice NumPy has made as well: ["As of NumPy 1.4.0 searchsorted works with real/complex arrays containing nan values. The enhanced sort order is documented in sort."](https://numpy.org/doc/stable/reference/generated/numpy.searchsorted.html) (And it would certainly handle my use case.)

The point about negative NaNs is certainly an oddity - since the concept of a negative NaN isn't technically defined - but agreed that this is still more sensible than the current behaviour.

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -50,7 +50,7 @@
 from jax.interpreters import pxla
 from jax import lax
 from jax._src import device_array
-from jax._src.lax.lax import _array_copy
+from jax._src.lax.lax import _array_copy, _float_to_int_for_sort
 from jax._src.ops import scatter
 from jax._src.util import (unzip2, prod as _prod, subvals, safe_zip, ceil_of_ratio,
                            canonicalize_axis as _canonicalize_axis, maybe_named_axis)
@@ -6456,7 +6456,11 @@ def _searchsorted(a, v, side):
   if len(a) == 0:
     return 0
   op = operator.le if side == 'left' else operator.lt
-
+  # TODO(jakevdp): handle NaNs correctly for complex. This will likely involve
+  # adding lexicographic sorting capabilities to the following.
+  a, v = _promote_dtypes(a, v)
+  if issubdtype(a.dtype, floating):
+    a, v = map(_float_to_int_for_sort, (a, v))
   def body_fun(i, state):
     low, high = state
     mid = (low + high) // 2


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2978,6 +2978,20 @@ def testSearchsorted(self, ashape, vshape, side, dtype):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+    {"testcase_name": f"_dtype={dtype.__name__}", "dtype": dtype}
+    for dtype in inexact_dtypes))
+  def testSearchsortedNans(self, dtype):
+    if np.issubdtype(dtype, np.complexfloating):
+      raise SkipTest("Known failure for complex inputs; see #9107")
+    sorted = jnp.array([-np.nan, -np.inf, -1, 0, 1, np.inf, np.nan], dtype=dtype)
+    self.assertArraysEqual(
+      jnp.searchsorted(sorted, sorted, side='left'),
+      jnp.arange(len(sorted)))
+    self.assertArraysEqual(
+      jnp.searchsorted(sorted, sorted, side='right'),
+      jnp.arange(1, 1 + len(sorted)))
+
   @parameterized.named_parameters(jtu.cases_from_list(
     {"testcase_name": "_x={}_bins={}_right={}_reverse={}".format(
       jtu.format_shape_dtype_string(xshape, dtype),