{"record":{"id":"fe9fb90eea9c91e6","repo":"jax-ml/jax","slug":"nan","errorCode":null,"errorMessage":"nan","messagePattern":"nan","errorType":"exception","errorClass":"InternalFloatingPointError","httpStatus":null,"severity":"error","filePath":"jax/_src/dispatch.py","lineNumber":311,"sourceCode":"\ndef check_special(name: str, bufs: Sequence[basearray.Array]) -> None:\n  if needs_check_special():\n    for buf in bufs:\n      _check_special(name, buf.dtype, buf)\n\n\ndef check_special_array(name: str, arr: array.ArrayImpl) -> array.ArrayImpl:\n  if needs_check_special():\n    if dtypes.issubdtype(arr.dtype, np.inexact):\n      for buf in arr._arrays:\n        _check_special(name, buf.dtype, buf)\n  return arr\n\n\ndef _check_special(name: str, dtype: np.dtype, buf: basearray.Array) -> None:\n  if dtypes.issubdtype(dtype, np.inexact):\n    if config.debug_nans.value and np.any(np.isnan(np.asarray(buf))):\n      raise InternalFloatingPointError(name, \"nan\")\n    if config.debug_infs.value and np.any(np.isinf(np.asarray(buf))):\n      raise InternalFloatingPointError(name, \"inf\")\n\ndef _device_put_reshard(x): return x\n\n\n@util.cache(max_size=2048, trace_context_in_key=False)\ndef _cached_logical_device_ids(\n    inp_device_list: xc.DeviceList,\n    target_device_list: xc.DeviceList\n) -> tuple[int, ...]:\n  device_to_index = {d: i for i, d in enumerate(target_device_list)}\n  return tuple(device_to_index[d] for d in inp_device_list)\n\n\ndef _different_device_order_reshard(\n    x: array.ArrayImpl, target_sharding: NamedSharding, copy: ArrayCopySemantics\n) -> array.ArrayImpl:","sourceCodeStart":293,"sourceCodeEnd":329,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/dispatch.py#L293-L329","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Inspect inputs and intermediate values with jax.debug.print or jnp.isnan(x).any() to locate the source of NaNs","Fix the numerical issue (clip values, stable log/softmax, lower learning rate, sanitize data)","Only enable jax_debug_nans while debugging — it disables caching and slows execution"],"exampleFix":"# before\nwith jax.debug_nans(True):\n    y = jnp.log(x)  # x contains negatives\n# after\nx = jnp.maximum(x, 1e-12)\ny = jnp.log(x)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\ndef has_nan(x): return bool(jnp.isnan(jnp.asarray(x, dtype=jnp.float32)).any())\nassert not any(has_nan(a) for a in jax.tree_util.tree_leaves(inputs)), 'NaN in inputs'","typeGuard":null,"tryCatchPattern":"try:\n    with jax.debug_nans(True):\n        out = f(x)\nexcept jax._src.dispatch.InternalFloatingPointError as e:\n    print('NaN produced by op:', e); raise","preventionTips":["Sanitize data pipelines: replace NaN/Inf at load time","Use stable primitives (logsumexp, softmax, clip)","Treat debug_nans as a temporary diagnostic, not production config"],"tags":["jax","nan","numerical-debugging"],"backgroundTag":"nan-in-training-output","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}