Repository: google/jax

Issue:
nanmedian (sometimes) incorrectly returns nan
`nanmedian` of a logged array containing `nan`s incorrectly returns `nan`

```python
import jax
import jax.numpy as np
import itertools

M = np.e # or w/e
a = np.array([0, M, 0])

# log of nans
la1 = np.log(np.array([np.nan, M, np.nan]))
la2 = np.log(np.where(a > 0, a, np.nan))
la3 = np.log(jax.ops.index_update(a, a <= 0, np.nan))

# nans inserted after log
la4 = np.array([np.nan, np.log(M), np.nan])
la5 = np.log(a); la5 = np.where(np.isinf(la5), np.nan, la5)
la6 = np.log(a); la6 = jax.ops.index_update(la5, np.isinf(la6), np.nan)

for pair in itertools.combinations((la1,la2,la3,la4,la5,la6), 2):
    assert np.array_equal(*pair, equal_nan=True) # all True

print(np.nanmedian(la1)) # nan
print(np.nanmedian(la2)) # nan
print(np.nanmedian(la3)) # nan
print(np.nanmedian(la4)) # log(M)
print(np.nanmedian(la5)) # log(M)
print(np.nanmedian(la6)) # log(M)
```

Other nan functions seem to work as expected, e.g.
```python
print(np.nanmean(la1)) # log(M)
print(np.nanmean(la4)) # log(M)
```


Discussion:
Thanks for the report! This one is pretty strange... it looks like for some reason the particular float representation of the nan value matters!
```python
print(la1.view('uint8'))
print(la4.view('uint8'))
# [255 255 255 255   0   0 128  63 255 255 255 255]
# [  0   0 192 127   0   0 128  63   0   0 192 127]
```
It probably has something to do with how these react to the sorting used under the hood in `median`.
indeed! interesting..

```
jax.lax.sort(la1) # DeviceArray([nan, nan,  1.], dtype=float32)
jax.lax.sort(la4) # DeviceArray([ 1., nan, nan], dtype=float32)
jax.lax.sort(np.hstack([la1,la4])) # DeviceArray([nan, nan,  1.,  1., nan, nan], dtype=float32)
```

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
@@ -5170,9 +5170,10 @@ def _quantile(a, q, axis, interpolation, keepdims, squash_nans):
     raise ValueError("q must be have rank <= 1, got shape {}".format(shape(q)))
 
   a_shape = shape(a)
-  a = lax.sort(a, dimension=axis)
 
   if squash_nans:
+    a = where(isnan(a), nan, a) # Ensure nans are positive so they sort to the end.
+    a = lax.sort(a, dimension=axis)
     counts = sum(logical_not(isnan(a)), axis=axis, dtype=q.dtype,
                  keepdims=keepdims)
     shape_after_reduction = counts.shape
@@ -5200,6 +5201,7 @@ def _quantile(a, q, axis, interpolation, keepdims, squash_nans):
     index[axis] = high
     high_value = a[tuple(index)]
   else:
+    a = lax.sort(a, dimension=axis)
     n = a_shape[axis]
     q = lax.mul(q, _constant_like(q, n - 1))
     low = lax.floor(q)


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -666,10 +666,13 @@ def rand(shape, dtype):
       return base_rand(shape, dtype)
 
     dims = _dims_of_shape(shape)
-    nan_flips = rng.rand(*dims) < 0.1
+    r = rng.rand(*dims)
+    nan_flips = r < 0.1
+    neg_nan_flips = r < 0.05
 
     vals = base_rand(shape, dtype)
     vals = np.where(nan_flips, np.array(np.nan, dtype=dtype), vals)
+    vals = np.where(neg_nan_flips, np.array(-np.nan, dtype=dtype), vals)
 
     return _cast_to_shape(np.asarray(vals, dtype=dtype), shape, dtype)