jax-ml/jax · error · RuntimeError

JAX tried to execute function {self.name}, but the no_execut

Error message

JAX tried to execute function {self.name}, but the no_execution config option is set

What it means

JAX supports a diagnostic mode (jax_no_execution / config.no_execution) that performs full lowering/setup but skips actually running computations. If a jitted function reaches __call__ with this flag on, it raises instead of executing — used to test compilation-only paths.

Source

Thrown at jax/_src/interpreters/pxla.py:402

      if len(token_buf) == 1:
        dispatch.runtime_tokens.set_token_result(eff, core.Token(token_buf[0]))
      else:
        token_devices = []
        for token in token_buf:
          assert len(token.sharding.device_set) == 1
          token_devices.append(token.sharding._device_assignment[0])
        s = NamedSharding(Mesh(token_devices, 'x'), P('x'))
        global_token_array = array.make_array_from_single_device_arrays(
            (0,), s, token_buf
        )
        dispatch.runtime_tokens.set_token_result(
            eff, core.Token(global_token_array)
        )

  @profiler.annotate_function
  def __call__(self, *args):
    if config.no_execution.value:
      raise RuntimeError(
      f"JAX tried to execute function {self.name}, but the no_execution config "
      "option is set")
    args = [x for i, x in enumerate(args) if i in self.kept_var_idx]
    if self.mut:
      args = [*args, *self.mut.in_mut]
    input_bufs = self.in_handler(args)
    with profiler.PGLEProfiler.trace(self.pgle_profiler):
      if (self.ordered_effects or self.has_unordered_effects
          or self.has_host_callbacks):
        input_bufs = self._add_tokens_to_inputs(input_bufs)
        results = self.xla_executable.execute_sharded(input_bufs, with_tokens=True)

        result_token_bufs = results.consume_with_handlers(
            [lambda xs: xs] * len(self.ordered_effects), strict=False)
        sharded_runtime_token = results.consume_token()
        self._handle_token_bufs(result_token_bufs, sharded_runtime_token)
      else:
        results = self.xla_executable.execute_sharded(input_bufs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Unset the environment variable (export JAX_NO_EXECUTION=false or unset JAX_NO_EXECUTION) or jax.config.update('no_execution', False)
If intentional for compile-only testing, don't call the function — stop at .compile()

Example fix

# before
# env: JAX_NO_EXECUTION=true
out = jitted_fn(x)

# after
# env: unset JAX_NO_EXECUTION
out = jitted_fn(x)

Defensive patterns

Strategy: validation

Validate before calling

import jax
if jax.config.no_execution.value:
    raise RuntimeError('unset JAX_NO_EXECUTION before calling functions')

Try / catch

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

Prevention

Don't export JAX_NO_EXECUTION globally in dev shells
In compile-only CI, stop at .compile() and never call the executable

When it happens

Trigger: Setting JAX_NO_EXECUTION=true (or jax.config.update('no_execution', True)) then calling a jitted/pjitted function. Compilation via .lower().compile() works; only invocation raises.

Common situations: CI environments verifying compilability; accidentally exported env var; stale shell config setting the variable globally.

Related errors

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