jax-ml/jax · error · InternalFloatingPointError

inf

Error message

inf

What it means

Raised when jax_debug_infs (debug_infs config) is enabled and a computed buffer contains infinity. Like the NaN check, it turns silent Inf values into an early InternalFloatingPointError at the op that produced them, at the cost of performance.

Source

Thrown at jax/_src/dispatch.py:313

  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:
  x._check_if_deleted()
  inp_sharding = x.sharding

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Locate the overflowing op via traceback/jax.debug.print and guard it (clip logits, use logsumexp, epsilon in denominators)
2. Switch to a wider dtype (float32) or enable loss scaling for mixed precision
3. Turn off jax_debug_infs once resolved, as it disables caching and adds overhead

Example fix

# before
p = jnp.exp(logits)  # overflow with debug_infs
# after
p = jax.nn.softmax(logits)  # uses stable logsumexp internally

import jax.numpy as jnp
def finite(x): return bool(jnp.isfinite(jnp.asarray(x, dtype=jnp.float32)).all())
assert all(finite(a) for a in jax.tree_util.tree_leaves(inputs))

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

When it happens

Trigger: Enabling jax_debug_infs while a computation overflows to ±Inf, e.g. exp(large), division by zero, or fp16 overflow; check_special inspects every op output buffer.

Common situations: Mixed-precision (bfloat16/fp16) training overflow; exp on large logits; division by near-zero denominators; enabled during numerical debugging sessions.

Related errors

• 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
• The names should be exclusive and should not intersect in `n

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