jax-ml/jax · error · FloatingPointError

invalid value ({e.ty}) encountered in ndtri.

Error message

invalid value ({e.ty}) encountered in ndtri.

What it means

After computing the inverse normal CDF, _ndtri runs dispatch.check_special to detect invalid values (NaN etc.) produced during the computation. If an internal floating point error of type e.ty (e.g. invalid) is found, it is re-raised as a user-facing FloatingPointError. This only happens for concrete (non-traced) arrays.

Source

Thrown at jax/_src/scipy/special.py:1597

  second_term_small_p = jnp.polyval(p2, 1 / z) / jnp.polyval(q2, 1 / z) / z
  second_term_otherwise = jnp.polyval(p1, 1 / z) / jnp.polyval(q1, 1 / z) / z
  x_for_small_p = first_term - second_term_small_p
  x_otherwise = first_term - second_term_otherwise

  x = jnp.where(sanitized_mcp > dtype(np.exp(-2.)),
                x_for_big_p,
                jnp.where(z >= dtype(8.0), x_for_small_p, x_otherwise))

  x = jnp.where(p > dtype(1. - np.exp(-2.)), x, -x)
  with config.debug_infs(False):
    infinity = jnp.full(shape, dtype(np.inf))
    x = jnp.where(
        p == dtype(0.0), -infinity, jnp.where(p == dtype(1.0), infinity, x))
  if not isinstance(x, core.Tracer):
    try:
      dispatch.check_special("ndtri", [x])
    except api_util.InternalFloatingPointError as e:
      raise FloatingPointError(
          f"invalid value ({e.ty}) encountered in ndtri.") from None
  return x


@partial(custom_derivatives.custom_jvp, nondiff_argnums=(1,))
def log_ndtr(x: ArrayLike, series_order: int = 3) -> Array:
  r"""Log Normal distribution function.

  JAX implementation of :obj:`scipy.special.log_ndtr`.

  For details of the Normal distribution function see `ndtr`.

  This function calculates :math:`\log(\mathrm{ndtr}(x))` by either calling
  :math:`\log(\mathrm{ndtr}(x))` or using an asymptotic series. Specifically:

  - For `x > upper_segment`, use the approximation `-ndtr(-x)` based on
    :math:`\log(1-x) \approx -x, x \ll 1`.
  - For `lower_segment < x <= upper_segment`, use the existing `ndtr` technique

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Sanitize p before calling: replace NaNs, e.g. p = jnp.where(jnp.isnan(p), 0.5, p), and clip to [0,1] with jnp.clip(p, 0., 1.)
2. Find the upstream NaN source: check p with jnp.isnan(p).any() and print before the call
3. Wrap in jax.jit if you must defer NaN checking, but treat that as masking, not fixing

Example fix

// before
q = jax.scipy.special.ndtri(p)  # p has NaNs
// after
p = jnp.where(jnp.isnan(p), 0.5, jnp.clip(p, 0.0, 1.0))
q = jax.scipy.special.ndtri(p)

p = jnp.asarray(p)
assert not jnp.isnan(p).any(), 'NaN in ndtri input'
p = jnp.clip(jnp.where(jnp.isnan(p), 0.5, p), 0.0, 1.0)

def no_nans(p):
    return not bool(jnp.isnan(p).any())

try:
    q = jax.scipy.special.ndtri(p)
except FloatingPointError:
    p = jnp.nan_to_num(p, nan=0.5)
    q = jax.scipy.special.ndtri(p)

When it happens

Trigger: Calling jax.scipy.special.ndtri(p) eagerly (outside jit) with p containing NaN, or values that cause the rational approximations to produce NaN (e.g. wrong-dtype edge cases or p far outside [0,1] due to upstream bugs).

Common situations: NaN probabilities leaking from a model (log of negative, 0/0) then passed to ndtri; eager debugging sessions where NaN checks are active but jit hides them; jax_debug_nan or debug configurations surfacing hidden NaNs.

Related errors

• nan
• betainc does not support complex-valued inputs.
• dawsn does not support complex-valued inputs.
• entr does not support complex-valued inputs.
• kl_div does not support complex-valued inputs.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/d2e24f4043433e3c. Report an issue: GitHub.