jax-ml/jax · error · ValueError

raise_if_error() should not be called within a traced contex

Error message

raise_if_error() should not be called within a traced context, such as within a jitted function.

What it means

raise_if_error() reads the global error-code buffer and converts the minimum (first) error code into a Python exception. It must run outside tracing: if the reduced error code is a Tracer (i.e. the call is executing under jit/pmap/vmap/grad), JAX raises this ValueError because exceptions cannot depend on traced values.

Source

Thrown at jax/_src/error_check.py:235

  """Raise an exception if the internal error state is set.

  This function should be called after a computation completes to check for any
  errors that were marked during execution via `set_error_if()`. If an error
  exists, it raises a `JaxValueError` with the corresponding error message.

  This function should not be called inside a traced function (e.g., inside
  :func:`jax.jit`). Doing so will raise a `ValueError`.

  Raises:
    JaxValueError: If the internal error state is set.
    ValueError: If called within a traced JAX function.
  """
  if _error_storage.ref is None:  # if not initialized, do nothing
    return

  error_code = _error_storage.ref[...].min()  # reduce to a single error code
  if isinstance(error_code, core.Tracer):
    raise ValueError(
        "raise_if_error() should not be called within a traced context, such as"
        " within a jitted function."
    )
  if error_code == np.uint32(_NO_ERROR):
    return
  _error_storage.ref[...] = lax.full(
      _error_storage.ref.shape,
      np.uint32(_NO_ERROR),
      sharding=_error_storage.ref.sharding,
  )  # clear the error code

  with _error_list_lock:
    if error_code < 0 or error_code >= len(_error_list):
      # Handle invalid error codes gracefully with a standard error message.
      # This can happen with corrupted AOT serialization data or negative
      # error codes that could lead to incorrect indexing.
      msg, traceback = _INVALID_ERROR_CODE_MSG, _INVALID_ERROR_CODE_TRACEBACK
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move raise_if_error() outside the jitted function: call jit(fn)(x) first, then raise_if_error()
Split the step: keep set_error_if inside jit, do raise_if_error in the eager host loop
For control flow dependent on errors inside jit, use jax.lax.cond on the error code value instead of Python exceptions

Example fix

# before
@jax.jit
def step(x):
    y = compute(x)
    set_error_if(y < 0)
    raise_if_error()  # ValueError: traced context
    return y

# after
@jax.jit
def step(x):
    y = compute(x)
    set_error_if(y < 0)
    return y

y = step(x)
raise_if_error()

Defensive patterns

Strategy: validation

Validate before calling

import jax
if not isinstance(jax.core.Tracer, type) or not any(isinstance(v, jax.core.Tracer) for v in []):
    pass
# practical guard: only call outside transformations
if not _under_trace():
    raise_if_error()

Type guard

def is_traced(x) -> bool:
    from jax.core import Tracer
    return isinstance(x, Tracer)

Prevention

Keep raise_if_error() in the eager host loop, never inside @jit/@vmap/@grad
Structure code as: jitted compute (with set_error_if) then raise_if_error on host
Use lax.cond on error codes for error-dependent traced control flow

When it happens

Trigger: Calling jax.error_check.raise_if_error() inside a jitted function, inside jax.grad, vmap, or any transformation that traces its body; also calling it in a function later wrapped in jit.

Common situations: Adding error checks by wrapping the whole training step (including the check) in @jax.jit; helper functions reused in both eager and traced contexts; copying example code that called raise_if_error at the end of a loss function that is later jitted.

Related errors

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