jax-ml/jax · error · NotImplementedError

Cannot write to the same ref in both cond and body of while

Error message

Cannot write to the same ref in both cond and body of while loop.

What it means

This is a check in the state-discharge rule for while_loop: if the cond jaxpr and the body jaxpr both contain WriteEffects targeting the same Ref input, discharging state is ambiguous and JAX raises NotImplementedError. Writing the same ref from both parts of the loop is not representable in the discharged form.

Source

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

    cond_jaxpr, body_jaxpr, cond_nconsts, body_nconsts):
  cond_consts_discharge, body_consts_discharge, carry_discharge = split_list(
      ctx.should_discharge, [cond_nconsts, body_nconsts])
  cond_consts, body_consts, carry = split_list(args, [cond_nconsts, body_nconsts])
  cond_consts_avals, body_consts_avals, carry_avals = split_list(ctx.in_avals,
                                                                 [cond_nconsts,
                                                                  body_nconsts])

  # Check if the same Ref is written to in both cond and body.
  cond_write_ids = {id(cond_consts_avals[effect.input])
    for effect in core.positional_effects(cond_jaxpr)
    if isinstance(effect, state.WriteEffect)}
  cond_has_writes = len(cond_write_ids) > 0
  body_write_ids = {id(body_consts_avals[effect.input])
    for effect in core.positional_effects(body_jaxpr)
    if isinstance(effect, state.WriteEffect)}
  write_to_both_ids = cond_write_ids & body_write_ids
  if write_to_both_ids:
    raise NotImplementedError(
        "Cannot write to the same ref in both cond and body of while loop.")

  cond_is_ref = [
      isinstance(aval, state.AbstractRef) and should
      for aval, should in zip(cond_consts_avals, cond_consts_discharge)
  ]
  remaining_cond_consts, cond_refs = partition_list(cond_is_ref, cond_consts)
  remaining_cond_const_avals, cond_ref_avals = partition_list(cond_is_ref,
                                                         cond_consts_avals)
  num_cond_refs = sum(cond_is_ref)
  num_remaining_cond_consts = cond_nconsts - num_cond_refs
  body_is_ref = [
      isinstance(aval, state.AbstractRef) and should
      for aval, should in zip(body_consts_avals, body_consts_discharge)
  ]
  remaining_body_consts, body_refs = partition_list(body_is_ref, body_consts)
  remaining_body_const_avals, body_ref_avals = partition_list(body_is_ref,
                                                         body_consts_avals)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Move the write out of the cond function so only the body writes the ref
If the cond genuinely needs the written value, read the ref in cond but perform the write only in body
Split the ref usage into two separate refs written by cond and body respectively, if semantics allow
If you believe discharge should support it, file an issue with jax-ml/jax including a minimal repro

Example fix

// before
def cond(c):
  ref.set(ref[...] + 1)
  return c < 10
def body(c):
  ref.set(ref[...] * 2)
  return c + 1
// after
def cond(c):
  return c < 10
def body(c):
  ref.set(ref[...] * 2)
  return c + 1

Defensive patterns

Strategy: validation

Validate before calling

# audit before compiling: ensure cond performs no ref writes
import jax
def cond_writes_refs(cond_fn, example_carry):
    jaxpr = jax.make_jaxpr(cond_fn)(example_carry)
    return any('Write' in type(x).__name__ or 'write' in str(x).lower() for x in jaxpr.effects)

Prevention

Design cond functions as pure predicates — reads only, never writes
Keep ref mutation confined to the loop body
Add unit tests for stateful loops under the discharge transformation

When it happens

Trigger: Using jax.State / Ref values (e.g. jax.experimental.io_effect, state.get/put, or jax.lax.withunding state APIs) where the same Ref is written in both the cond function and the body function of lax.while_loop, and then the discharge transformation runs (e.g. under state discharge transformations like when the function is lowered or io_effect is discharged).

Common situations: Experimental stateful JAX code (io_effect, pallas or state API) where a helper that writes a ref is called from both cond and body; typically surfaced when jit-compiling stateful while loops on newer JAX versions.

Related errors

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