jax-ml/jax · error · InternalFloatingPointError

nan

Error message

nan

What it means

Raised when jax_check_nans/debug_nans is enabled and an intermediate or output array contains NaN. JAX intentionally converts silent NaNs into an InternalFloatingPointError so numerical bugs surface at the op that produced them rather than propagating downstream.

Source

Thrown at jax/_src/dispatch.py:311

def check_special(name: str, bufs: Sequence[basearray.Array]) -> None:
  if needs_check_special():
    for buf in bufs:
      _check_special(name, buf.dtype, buf)


def check_special_array(name: str, arr: array.ArrayImpl) -> array.ArrayImpl:
  if needs_check_special():
    if dtypes.issubdtype(arr.dtype, np.inexact):
      for buf in arr._arrays:
        _check_special(name, buf.dtype, buf)
  return arr


def _check_special(name: str, dtype: np.dtype, buf: basearray.Array) -> None:
  if dtypes.issubdtype(dtype, np.inexact):
    if config.debug_nans.value and np.any(np.isnan(np.asarray(buf))):
      raise InternalFloatingPointError(name, "nan")
    if config.debug_infs.value and np.any(np.isinf(np.asarray(buf))):
      raise InternalFloatingPointError(name, "inf")

def _device_put_reshard(x): return x


@util.cache(max_size=2048, trace_context_in_key=False)
def _cached_logical_device_ids(
    inp_device_list: xc.DeviceList,
    target_device_list: xc.DeviceList
) -> tuple[int, ...]:
  device_to_index = {d: i for i, d in enumerate(target_device_list)}
  return tuple(device_to_index[d] for d in inp_device_list)


def _different_device_order_reshard(
    x: array.ArrayImpl, target_sharding: NamedSharding, copy: ArrayCopySemantics
) -> array.ArrayImpl:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect inputs and intermediate values with jax.debug.print or jnp.isnan(x).any() to locate the source of NaNs
2. Fix the numerical issue (clip values, stable log/softmax, lower learning rate, sanitize data)
3. Only enable jax_debug_nans while debugging — it disables caching and slows execution

Example fix

# before
with jax.debug_nans(True):
    y = jnp.log(x)  # x contains negatives
# after
x = jnp.maximum(x, 1e-12)
y = jnp.log(x)

import jax.numpy as jnp
def has_nan(x): return bool(jnp.isnan(jnp.asarray(x, dtype=jnp.float32)).any())
assert not any(has_nan(a) for a in jax.tree_util.tree_leaves(inputs)), 'NaN in inputs'

try:
    with jax.debug_nans(True):
        out = f(x)
except jax._src.dispatch.InternalFloatingPointError as e:
    print('NaN produced by op:', e); raise

When it happens

Trigger: Enabling jax_debug_nans (or debug_nans config) while running any computation whose buffers contain NaN, e.g. 0/0, log of negatives, unstable learning rates. check_special runs on op outputs during dispatch.

Common situations: Debugging a diverging training run with jax_debug_nans=True; NaNs from mixed precision underflow; NaNs in input data going undetected until the flag is on.

Related errors

• invalid value ({e.ty}) encountered in ndtri.
• Invalid value ({e.ty}) encountered in sharded computation.
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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