Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

jax_debug_nans

See original GitHub issue

The current implementation of quantile function generates NaNs that are detected when jax_debug_nans is activated.

from jax.config import config
config.update("jax_debug_nans", True)
import jax.numpy as jnp
jnp.quantile(jnp.ones((3, 3)), 0.5)

from jax.config import config
config.update("jax_debug_nans", True)
import jax.numpy as jnp
jnp.where(1, float('nan'), 1)

Traceback

  File "/home/esac/projects/venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5560, in percentile
    return quantile(a, q, axis=axis, out=out, overwrite_input=overwrite_input,
  File "/home/esac/projects/venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5370, in quantile
    return _quantile(a, q, axis, interpolation, keepdims, False)
  File "/home/esac/projects/venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5436, in _quantile
    a = where(any(isnan(a), axis=axis, keepdims=True), nan, a)
  File "/home/esac/projects/venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1686, in where
    return _where(condition, x, y)
  File "/home/esac/projects/venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1669, in _where
    x, y = _promote_dtypes(x, y)
  File "/home/esac/projects/venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 268, in _promote_dtypes
    return [lax._convert_element_type(x, to_dtype, weak_type) for x in args]
  File "/home/esac/projects/venv/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 268, in <listcomp>
    return [lax._convert_element_type(x, to_dtype, weak_type) for x in args]
FloatingPointError: invalid value (nan) encountered in convert_element_type
python-BaseException
Invalid value encountered in the output of a jit function. Calling the de-optimized version.