{"record":{"id":"6499e6257b569168","repo":"jax-ml/jax","slug":"raise-if-error-should-not-be-called-within-a-tra","errorCode":null,"errorMessage":"raise_if_error() should not be called within a traced context, such as within a jitted function.","messagePattern":"raise_if_error\\(\\) should not be called within a traced context, such as within a jitted function\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/error_check.py","lineNumber":235,"sourceCode":"  \"\"\"Raise an exception if the internal error state is set.\n\n  This function should be called after a computation completes to check for any\n  errors that were marked during execution via `set_error_if()`. If an error\n  exists, it raises a `JaxValueError` with the corresponding error message.\n\n  This function should not be called inside a traced function (e.g., inside\n  :func:`jax.jit`). Doing so will raise a `ValueError`.\n\n  Raises:\n    JaxValueError: If the internal error state is set.\n    ValueError: If called within a traced JAX function.\n  \"\"\"\n  if _error_storage.ref is None:  # if not initialized, do nothing\n    return\n\n  error_code = _error_storage.ref[...].min()  # reduce to a single error code\n  if isinstance(error_code, core.Tracer):\n    raise ValueError(\n        \"raise_if_error() should not be called within a traced context, such as\"\n        \" within a jitted function.\"\n    )\n  if error_code == np.uint32(_NO_ERROR):\n    return\n  _error_storage.ref[...] = lax.full(\n      _error_storage.ref.shape,\n      np.uint32(_NO_ERROR),\n      sharding=_error_storage.ref.sharding,\n  )  # clear the error code\n\n  with _error_list_lock:\n    if error_code < 0 or error_code >= len(_error_list):\n      # Handle invalid error codes gracefully with a standard error message.\n      # This can happen with corrupted AOT serialization data or negative\n      # error codes that could lead to incorrect indexing.\n      msg, traceback = _INVALID_ERROR_CODE_MSG, _INVALID_ERROR_CODE_TRACEBACK\n    else:","sourceCodeStart":217,"sourceCodeEnd":253,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/error_check.py#L217-L253","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\n@jax.jit\ndef step(x):\n    y = compute(x)\n    set_error_if(y < 0)\n    raise_if_error()  # ValueError: traced context\n    return y\n\n# after\n@jax.jit\ndef step(x):\n    y = compute(x)\n    set_error_if(y < 0)\n    return y\n\ny = step(x)\nraise_if_error()","handlingStrategy":"validation","validationCode":"import jax\nif not isinstance(jax.core.Tracer, type) or not any(isinstance(v, jax.core.Tracer) for v in []):\n    pass\n# practical guard: only call outside transformations\nif not _under_trace():\n    raise_if_error()","typeGuard":"def is_traced(x) -> bool:\n    from jax.core import Tracer\n    return isinstance(x, Tracer)","tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","error-checking","jit","tracer"],"backgroundTag":"tracer-in-eager-context","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}