{"record":{"id":"cf7a019641ec76f6","repo":"jax-ml/jax","slug":"effects-not-supported-in-while","errorCode":null,"errorMessage":"Effects not supported in `while`: {}","messagePattern":"Effects not supported in `while`: (.+?)","errorType":"error_code","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":1719,"sourceCode":"  if len(body_out_avals) != len(init_aval):\n    _check_carry_type('while_loop body', body_fun, init_aval, body_out_avals)\n    assert False, \"shouldn't get here\"\n\n  init_val_flat, changed = init_val_flat.map3(\n      list(init_aval), body_out_avals,\n      _promote_weak_typed_input).unzip2()\n  if any(changed):\n    init_aval = init_val_flat.map(core.typeof)\n    cond_jaxpr, body_jaxpr, body_out_avals = _create_jaxpr(init_aval)\n\n  cond_jaxpr, cond_consts = pe.separate_consts(cond_jaxpr)\n  body_jaxpr, body_consts = pe.separate_consts(body_jaxpr)\n  _check_carry_type('while_loop body', body_fun, init_aval, body_out_avals)\n\n  joined_effects = core.join_effects(cond_jaxpr.effects, body_jaxpr.effects)\n  disallowed_effects = effects.control_flow_allowed_effects.filter_not_in(joined_effects)\n  if disallowed_effects:\n    raise NotImplementedError(\n        f'Effects not supported in `while`: {disallowed_effects}')\n\n  # If the body forwards an input carry to an output carry, *and* it's not used\n  # by the cond fun, it can be moved to be a body const. Doing so can lead to\n  # efficiency wins: if e.g. we vmap the loop with a batched predicate, we batch\n  # the carry too, but not the body consts.\n  body_fwd = pe._jaxpr_forwarding(body_jaxpr)\n  carry_nofwd = [len(body_consts) + i != f for i, f in enumerate(body_fwd)]\n  cond_jaxpr_, keep_cond = pe.dce_jaxpr(\n      cond_jaxpr, [True], [True] * len(cond_consts) + carry_nofwd)\n  _, keep_cond_carry = split_list(keep_cond, [len(cond_consts)])\n  move_to_const = _map(operator.not_, keep_cond_carry)\n\n  init_vals = list(init_val_flat)\n  new_body_consts: list[Any] = []\n  if any(move_to_const):\n    cond_jaxpr = cond_jaxpr_\n    body_jaxpr = pe.prune_closed_jaxpr_outputs(","sourceCodeStart":1701,"sourceCodeEnd":1737,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L1701-L1737","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Move the effectful operation out of the loop (compute before/after, or hoist to a host-side wrapper)","Replace jax.debug.print with inspection via jax.debug.inspect_value or run with JAX disabled JIT (config.disable_jit) where effects execute eagerly","Register/allow the effect if it's your own custom effect by adding it to the allowed-effects set for control flow","If it's a JAX-version regression, pin/upgrade JAX and report at https://github.com/jax-ml/jax/issues"],"exampleFix":"// before\ndef body(c):\n    jax.debug.print('c={}', c)  # ordered IO effect\n    return c + 1\njax.lax.while_loop(cond, body, 0)\n\n// after\ndef body(c):\n    return c + 1\nc_final = jax.lax.while_loop(cond, body, 0)\nprint('final', c_final)  # observe outside the loop","handlingStrategy":"validation","validationCode":"import jax\nallowed = jax._src.effects.control_flow_allowed_effects\ncond_eff = jax.make_jaxpr(cond_fun)(init_val).effects\nbody_eff = jax.make_jaxpr(body_fun)(init_val).effects\nassert (cond_eff | body_eff).issubset(allowed), f'disallowed: {cond_eff | body_eff}'","typeGuard":null,"tryCatchPattern":"try:\n    jax.lax.while_loop(cond, body, init)\nexcept NotImplementedError as e:\n    if 'Effects not supported' in str(e):\n        # hoist effectful op out of the loop, or debug with disable_jit\n        raise","preventionTips":["Never put jax.debug.print or callbacks inside traced loop bodies in production code","Debug loops with config.enable_checks / disable_jit runs instead of inline prints","Review effect allow-list changes in JAX release notes when upgrading"],"tags":["jax","effects","while-loop","debug-print","not-implemented"],"backgroundTag":"effectful-operation-inside-compiled-loop","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}