jax-ml/jax · error · RuntimeError

re-tracing function {fun} for `jit`, but 'no_tracing' is set

Error message

re-tracing function {fun} for `jit`, but 'no_tracing' is set

What it means

JAX has a diagnostic mode (jax_config no_tracing, often enabled via JAX_NO_TRACING=true) that forbids tracing: it exists to verify code is fully covered by caches/exports. If a jitted function misses the cache and needs re-tracing while this flag is on, tracing aborts with this RuntimeError.

Source

Thrown at jax/_src/interpreters/partial_eval.py:2058

    debug_info = debug_info._replace(arg_names=lo_arg_names)
  if debug_info.result_paths is not None:
    lo_result_paths = tuple(
    path for aval, path in zip(hi_jaxpr.out_avals, debug_info.result_paths)
        for _ in aval.lo_ty())
    debug_info = debug_info._replace(result_paths=lo_result_paths)
  return debug_info

def trace_to_jaxpr_nocache(
    fun: Callable,
    in_avals: ft.FlatTree,  # (args, kwargs) pair
    debug_info: core.DebugInfo,
    # TODO: let's just make a `trace_to_jaxpr_ft` function for this
    fun_takes_flat_tree_arg=False,
    fun_returns_flat_tree=False,
    requires_low=False,
) -> tuple[Jaxpr, ft.FlatTree]:
  if config.no_tracing.value:
    raise RuntimeError(f"re-tracing function {fun} for "
                       "`jit`, but 'no_tracing' is set")
  test_event("trace_to_jaxpr")
  config.enable_checks.value and debug_info.assert_arg_names(len(in_avals))
  parent_trace = core.trace_ctx.trace
  trace = DynamicJaxprTrace(debug_info, parent_trace=parent_trace,
                            lower=requires_low)
  # Name stack and the traceback scope are reset because the metadata on jaxpr
  # equations should be rooted at the enclosing jaxpr and not contain any
  # context from the callsite. Otherwise metadata from one caller would bleed
  # into metadata from a different caller if we, e.g., inline.
  with (core.ensure_no_leaks(trace), source_info_util.reset_name_stack(),
        TracebackScope()):
    source_info = source_info_util.current()
    if requires_low:
      if debug_info.arg_names is not None:
        debug_info = debug_info._replace(arg_names=tuple(
            name for aval, name in zip(in_avals, debug_info.arg_names)
            for _ in aval.lo_ty()))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pre-warm the cache: call the function once with all expected shapes/dtypes before enabling no_tracing, or use AOT (jax.export / .lower(...).compile()) artifacts
Unset the flag (remove JAX_NO_TRACING / jax.config.update('no_tracing', False)) if dynamic tracing is expected
Ensure argument shapes/dtypes/static keys match exactly what was traced, so no cache-miss re-trace occurs

Example fix

# before
# JAX_NO_TRACING=true in environment, first call of jitted fn
result = jitted_fn(x)

# after
# warm cache in setup, before enabling no_tracing:
result = jitted_fn(x_example)
# then enable jax.config.update('no_tracing', True)

Defensive patterns

Strategy: validation

Validate before calling

import jax
assert not jax.config.no_tracing.value or cache_is_warm(fn, args), \
    'no_tracing set but function not yet compiled'

Try / catch

try:
    out = jitted_fn(x)
except RuntimeError as e:
    if 'no_tracing' in str(e):
        jax.config.update('no_tracing', False)
        out = jitted_fn(x)

Prevention

Warm caches or use AOT exports before enabling no_tracing
Keep argument signatures stable in no_tracing environments

When it happens

Trigger: Setting config.no_tracing (e.g. export JAX_NO_TRACING=true) and then calling a jitted function that is not yet compiled/cached — the cache miss forces trace_to_jaxpr, which refuses. Also triggered by anything invalidating the cache (new argument shapes/dtypes, first call).

Common situations: CI or production environments that set JAX_NO_TRACING to enforce 'no surprise compilation'; first-call after process start; new input signatures after enabling the flag.

Related errors

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