{"record":{"id":"82251569022e5743","repo":"jax-ml/jax","slug":"invalid-value-e-ty-encountered-in-sharded-comp","errorCode":null,"errorMessage":"Invalid value ({e.ty}) encountered in sharded computation.","messagePattern":"Invalid value \\((.+?)\\) encountered in sharded computation\\.","errorType":"exception","errorClass":"FloatingPointError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1352,"sourceCode":"      return api.jit(fn, out_shardings=NamedSharding(mesh, dst_pspec))(x)\n    return api.jit(fn)(x)\n\ndef _match(mesh, check_vma, manual_axes, src_pspec, dst_pspec, x):\n  return shard_map(_rem_singleton, mesh=mesh, in_specs=src_pspec,\n                   out_specs=dst_pspec, check_vma=check_vma,\n                   axis_names=manual_axes)(x)\n\ndef _rem_singleton(x): return lax.squeeze(x, [0])\ndef _add_singleton(x): return lax.expand_dims(x, [0])\n\ndef _maybe_check_special(outs):\n  if not config.debug_nans.value and not config.debug_infs.value: return\n  bufs = [s.data for leaf in tree_leaves(outs)\n          for s in getattr(leaf, 'addressable_shards', [])]\n  try:\n    dispatch.check_special('shard_map', bufs)\n  except api_util.InternalFloatingPointError as e:\n    raise FloatingPointError(f'Invalid value ({e.ty}) encountered in sharded computation.') from None\n\nclass ShardMapTrace(core.Trace):\n  __slots__ = (\"mesh\", \"manual_axes\", \"check\", \"amesh\")\n\n  mesh: Mesh  # outer concrete or abstract mesh\n  manual_axes: frozenset[AxisName]\n  check: bool\n\n  def __init__(self, mesh, manual_axes, check):\n    super().__init__()\n    self.mesh = mesh\n    self.manual_axes = manual_axes\n    self.check = check\n    self.amesh = mesh.abstract_mesh\n\n  def to_val_mat_pair(self, val):\n    if isinstance(val, ShardMapTracer):\n      return val.val, val.mat","sourceCodeStart":1334,"sourceCodeEnd":1370,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1334-L1370","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Locate the NaN/Inf source with jax.debug.print or nan debugging inside the body","Fix numerics (clipping, epsilon, dtype promotion, safe division)","Disable debug flags if they were enabled unintentionally"],"exampleFix":"# before\njax.config.update(\"jax_debug_nans\", True)  # surfaces FloatingPointError\n# after (after fixing e.g. log of negative input)\nx = jnp.clip(jnp.log(x), a_min=-1e30)","handlingStrategy":"try-catch","validationCode":null,"typeGuard":null,"tryCatchPattern":"import jax\ntry:\n    y = f(x)\nexcept jax.exceptions.FloatingPointError as e:\n    if 'sharded computation' in str(e): dump intermediate shards with jax.debug.print and halt training\n    else: raise","preventionTips":["Keep debug_nans/inf flags on in dev builds to catch divergence early","Add gradient/value clipping in distributed training loops"],"tags":["shard-map","nan","inf","debug-flags","jax"],"backgroundTag":"nan-in-debug-mode","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}