jax-ml/jax · error · JaxRuntimeError

{str(exc)}

Error message

{str(exc)}

What it means

This is the user-facing re-raise when a checkify-instrumented function fails at runtime: the deferred error collected by checkify.check is converted to JaxRuntimeError with the original formatted message.

Source

Thrown at jax/_src/checkify.py:488

core.pp_eqn_rules[check_p] = _pp_check

# TODO(lenamartens): inherit from Exception instead of ValueError.
class JaxRuntimeError(ValueError):
  pass

@check_p.def_impl
def check_impl(*args, err_tree, debug):
  if debug:
    # NOOP (check will only trigger when discharged)
    return []
  error = tree_unflatten(err_tree, args)
  exc = error.get_exception()
  if exc:
    filtered_tb = traceback_util.filter_traceback(
        exc.traceback_info.as_python_traceback())
    exc.with_traceback(filtered_tb)
    raise JaxRuntimeError(str(exc)) from exc
  return []

@check_p.def_effectful_abstract_eval
def check_abstract_eval(*args, err_tree, debug):
  del debug
  return [], set(tree_unflatten(err_tree, args)._pred.keys())

# TODO(lenamartens) add in-depth error explanation to link to in module docs.
functionalization_error = ValueError(
    'Cannot abstractly evaluate a checkify.check which was not'
    ' functionalized. This probably means you tried to stage'
    ' (jit/scan/pmap/...) a `check` without functionalizing it'
    ' through `checkify.checkify`.'
    )

def check_lowering_rule(ctx, *args, err_tree, debug):
  if debug:
    # NOOP (check will only trigger when discharged)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the message: it names which check failed and its formatted arguments
2. Fix the underlying logic so the check predicate passes
3. If NaNs are the cause, gradient-clip, tune learning rate, or sanitize inputs; use jax.debug_nans to locate them

Example fix

# before
checked_fn = checkify.checkify(f)
err, out = checked_fn(x)  # later raises JaxRuntimeError('nan check failed')
# after
def f(x):
  checkify.check(~jnp.isnan(x).any(), 'nan in x')
  return jnp.where(jnp.isnan(x), 0., x)

err, out = checked_fn(x)  # functional API avoids the raise
if err:
    print(err.get())  # inspect before it ever raises
    handle(out)

When it happens

Trigger: Calling a function wrapped with checkify.checkify(...) (and checkified) whose check/debug_check predicate evaluated False during execution.

Common situations: Numeric assertions like NaN checks, index bounds, or custom invariants failing inside jit-compiled code; the whole point of checkify — surfacing logical errors from compiled code.

Related errors

• Checkify does not support batched while-loops (checkify-of-v
• Mesh must be provided for shard_map with checkify.
• Unsupported aval type: {type(v)}
• {prim_name} takes a scalar pred as argument, got {pred}
• Formatting arguments to checkify.check need to be PyTrees of

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