jax-ml/jax · error · NotImplementedError

Effects not supported in `while`: {}

Error message

Effects not supported in `while`: {}

What it means

while_loop is a compiled control-flow primitive, so only effects in effects.control_flow_allowed_effects may occur inside cond_fun/body_fun. If tracing reveals other effects (e.g. host callbacks, ordered IO like debug prints, random state in older JAX), JAX raises NotImplementedError at the public API level rather than producing ill-defined execution semantics.

Source

Thrown at jax/_src/lax/control_flow/loops.py:1719

  if len(body_out_avals) != len(init_aval):
    _check_carry_type('while_loop body', body_fun, init_aval, body_out_avals)
    assert False, "shouldn't get here"

  init_val_flat, changed = init_val_flat.map3(
      list(init_aval), body_out_avals,
      _promote_weak_typed_input).unzip2()
  if any(changed):
    init_aval = init_val_flat.map(core.typeof)
    cond_jaxpr, body_jaxpr, body_out_avals = _create_jaxpr(init_aval)

  cond_jaxpr, cond_consts = pe.separate_consts(cond_jaxpr)
  body_jaxpr, body_consts = pe.separate_consts(body_jaxpr)
  _check_carry_type('while_loop body', body_fun, init_aval, body_out_avals)

  joined_effects = core.join_effects(cond_jaxpr.effects, body_jaxpr.effects)
  disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(joined_effects)
  if disallowed_effects:
    raise NotImplementedError(
        f'Effects not supported in `while`: {disallowed_effects}')

  # If the body forwards an input carry to an output carry, *and* it's not used
  # by the cond fun, it can be moved to be a body const. Doing so can lead to
  # efficiency wins: if e.g. we vmap the loop with a batched predicate, we batch
  # the carry too, but not the body consts.
  body_fwd = pe._jaxpr_forwarding(body_jaxpr)
  carry_nofwd = [len(body_consts) + i != f for i, f in enumerate(body_fwd)]
  cond_jaxpr_, keep_cond = pe.dce_jaxpr(
      cond_jaxpr, [True], [True] * len(cond_consts) + carry_nofwd)
  _, keep_cond_carry = split_list(keep_cond, [len(cond_consts)])
  move_to_const = _map(operator.not_, keep_cond_carry)

  init_vals = list(init_val_flat)
  new_body_consts: list[Any] = []
  if any(move_to_const):
    cond_jaxpr = cond_jaxpr_
    body_jaxpr = pe.prune_closed_jaxpr_outputs(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move the effectful operation out of the loop (compute before/after, or hoist to a host-side wrapper)
2. Replace jax.debug.print with inspection via jax.debug.inspect_value or run with JAX disabled JIT (config.disable_jit) where effects execute eagerly
3. Register/allow the effect if it's your own custom effect by adding it to the allowed-effects set for control flow
4. If it's a JAX-version regression, pin/upgrade JAX and report at https://github.com/jax-ml/jax/issues

Example fix

// before
def body(c):
    jax.debug.print('c={}', c)  # ordered IO effect
    return c + 1
jax.lax.while_loop(cond, body, 0)

// after
def body(c):
    return c + 1
c_final = jax.lax.while_loop(cond, body, 0)
print('final', c_final)  # observe outside the loop

import jax
allowed = jax._src.effects.control_flow_allowed_effects
cond_eff = jax.make_jaxpr(cond_fun)(init_val).effects
body_eff = jax.make_jaxpr(body_fun)(init_val).effects
assert (cond_eff | body_eff).issubset(allowed), f'disallowed: {cond_eff | body_eff}'

try:
    jax.lax.while_loop(cond, body, init)
except NotImplementedError as e:
    if 'Effects not supported' in str(e):
        # hoist effectful op out of the loop, or debug with disable_jit
        raise

When it happens

Trigger: Calling debug callbacks (jax.debug.print, jax.debug.callback), effectful random generators, or custom primitives declaring non-allowed effects inside while_loop's cond_fun or body_fun; also triggered when transformations re-trace the loop with newly-joined effects.

Common situations: Adding debug printing inside a while_loop for troubleshooting; upgrading JAX where an effect (e.g. new effect class for io/rand) moved out of the allowed set; using custom primitives with declared effects in training loops; interaction with remat/pjit re-tracing.

Related errors

• The following ordered effects are not supported for more tha
• Ordered IO effects not supported in vmap.
• Checkify does not support batched while-loops (checkify-of-v
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.

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