Repository: jax-ml/jax
Issue category: Bug Report

Issue:
median FloatingPointError: invalid value (nan) encountered in jit(convert_element_type)
### Description

Hello, 
I got this error in jnp.median when I set JAX_DISABLE_JIT=True and JAX_DEBUG_NANS=True.
```
Traceback (most recent call last):
  File "/data1/home/hhu17/zyl/PINE/H2+/3/test.py", line 29, in <module>
    c = jnp.median(b)
        ^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/numpy/reductions.py", line 2517, in median
    return quantile(a, 0.5, axis=axis, out=out, overwrite_input=overwrite_input,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/numpy/reductions.py", line 2172, in quantile
    return _quantile(lax_internal.asarray(a), lax_internal.asarray(q), axis, method, keepdims, False)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/numpy/reductions.py", line 2302, in _quantile
    a = _where(any(lax_internal._isnan(a), axis=axis, keepdims=True), np.nan, a)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/numpy/util.py", line 426, in _where
    x, y = promote_dtypes(x, y)
           ^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/numpy/util.py", line 259, in promote_dtypes
    return [lax._convert_element_type(x, to_dtype, weak_type) for x in args]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/numpy/util.py", line 259, in <listcomp>
    return [lax._convert_element_type(x, to_dtype, weak_type) for x in args]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 587, in _convert_element_type
    return convert_element_type_p.bind(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/lax/lax.py", line 2981, in _convert_element_type_bind
    operand = core.Primitive.bind(convert_element_type_p, operand,
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/core.py", line 438, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/core.py", line 442, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/core.py", line 955, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data1/home/hhu17/.conda/envs/jax/lib/python3.11/site-packages/jax/_src/dispatch.py", line 91, in apply_primitive
    outs = fun(*args)
           ^^^^^^^^^^
FloatingPointError: invalid value (nan) encountered in jit(convert_element_type). Because jax_config.debug_nans.value and/or config.jax_debug_infs is set, the de-optimized function (i.e., the function as if the `jit` decorator were removed) was called in an attempt to get a more precise error message. However, the de-optimized function did not produce invalid values during its execution. This behavior can result from `jit` optimizations causing the invalid value to be produced. It may also arise from having nan/inf constants as outputs, like `jax.jit(lambda ...: jax.numpy.nan)(...)`.
```

Following is the minimal code to reproduce the error.
```
import jax.numpy as jnp
import jax

key = jax.random.PRNGKey(12)
a = jax.random.normal(key, 128)
b = jnp.array(a)
c = jnp.median(b)
jit_median = jax.jit(jnp.median)
c = jit_median(b)
print(c)
```

Any help would be greatly appreciated!

### System info (python version, jaxlib version, accelerator, etc.)

jax:    0.4.35
jaxlib: 0.4.34
numpy:  2.1.1
python: 3.11.9 (main, Apr 19 2024, 16:48:06) [GCC 11.2.0]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='login4', release='3.10.0-957.el7.x86_64', version='#1 SMP Mon Dec 7 11:30:56 UTC 2020', machine='x86_64')


Hints:
Looks like it's coming from the NaN introduced on this line:

https://github.com/jax-ml/jax/blob/4b4fb9dae9eb7e2740d70de5b4a610f979530382/jax/_src/numpy/reductions.py#L2363

@jakevdp Can I tag you here since you wrote the implementation for _quantile?

Patch:
diff --git a/jax/_src/numpy/reductions.py b/jax/_src/numpy/reductions.py
index fa8d73361e2b..be1e55675079 100644
--- a/jax/_src/numpy/reductions.py
+++ b/jax/_src/numpy/reductions.py
@@ -2360,7 +2360,8 @@ def _quantile(a: Array, q: Array, axis: int | tuple[int, ...] | None,
     index[axis] = high
     high_value = a[tuple(index)]
   else:
-    a = _where(any(lax_internal._isnan(a), axis=axis, keepdims=True), np.nan, a)
+    with jax.debug_nans(False):
+      a = _where(any(lax_internal._isnan(a), axis=axis, keepdims=True), np.nan, a)
     a = lax.sort(a, dimension=axis)
     n = lax.convert_element_type(a_shape[axis], lax_internal._dtype(q))
     q = lax.mul(q, n - 1)


Test patch: