jax-ml/jax · error · FloatingPointError

Invalid value ({e.ty}) encountered in sharded computation.

Error message

Invalid value ({e.ty}) encountered in sharded computation.

What it means

When jax_debug_nans or jax_debug_infs is enabled, shard_map checks its output buffers for NaN/Inf and re-raises as a FloatingPointError identifying which invalid value type was produced inside the sharded computation.

Source

Thrown at jax/_src/shard_map.py:1352

      return api.jit(fn, out_shardings=NamedSharding(mesh, dst_pspec))(x)
    return api.jit(fn)(x)

def _match(mesh, check_vma, manual_axes, src_pspec, dst_pspec, x):
  return shard_map(_rem_singleton, mesh=mesh, in_specs=src_pspec,
                   out_specs=dst_pspec, check_vma=check_vma,
                   axis_names=manual_axes)(x)

def _rem_singleton(x): return lax.squeeze(x, [0])
def _add_singleton(x): return lax.expand_dims(x, [0])

def _maybe_check_special(outs):
  if not config.debug_nans.value and not config.debug_infs.value: return
  bufs = [s.data for leaf in tree_leaves(outs)
          for s in getattr(leaf, 'addressable_shards', [])]
  try:
    dispatch.check_special('shard_map', bufs)
  except api_util.InternalFloatingPointError as e:
    raise FloatingPointError(f'Invalid value ({e.ty}) encountered in sharded computation.') from None

class ShardMapTrace(core.Trace):
  __slots__ = ("mesh", "manual_axes", "check", "amesh")

  mesh: Mesh  # outer concrete or abstract mesh
  manual_axes: frozenset[AxisName]
  check: bool

  def __init__(self, mesh, manual_axes, check):
    super().__init__()
    self.mesh = mesh
    self.manual_axes = manual_axes
    self.check = check
    self.amesh = mesh.abstract_mesh

  def to_val_mat_pair(self, val):
    if isinstance(val, ShardMapTracer):
      return val.val, val.mat

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Locate the NaN/Inf source with jax.debug.print or nan debugging inside the body
2. Fix numerics (clipping, epsilon, dtype promotion, safe division)
3. Disable debug flags if they were enabled unintentionally

Example fix

# before
jax.config.update("jax_debug_nans", True)  # surfaces FloatingPointError
# after (after fixing e.g. log of negative input)
x = jnp.clip(jnp.log(x), a_min=-1e30)

import jax
try:
    y = f(x)
except jax.exceptions.FloatingPointError as e:
    if 'sharded computation' in str(e): dump intermediate shards with jax.debug.print and halt training
    else: raise

When it happens

Trigger: Enabling JAX debugging flags (jax.config.update('jax_debug_nans', True) or debug_infs) while the mapped function produces NaNs or Infs on any shard.

Common situations: Debugging divergence in distributed training; the flags are on in CI/dev and off in prod, so the error appears only in some environments; a bad shard (uneven data) producing overflow.

Related errors

• Mesh must be provided for shard_map with checkify.
• Unsupported aval type: {type(v)}
• Primitive {prim_name} requires varying manual axes to match,
• nan
• inf

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