{"record":{"id":"995529323d9be4b2","repo":"jax-ml/jax","slug":"cannot-write-to-the-same-ref-in-both-cond-and-body","errorCode":null,"errorMessage":"Cannot write to the same ref in both cond and body of while loop.","messagePattern":"Cannot write to the same ref in both cond and body of while loop\\.","errorType":"error_code","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2267,"sourceCode":"    cond_jaxpr, body_jaxpr, cond_nconsts, body_nconsts):\n  cond_consts_discharge, body_consts_discharge, carry_discharge = split_list(\n      ctx.should_discharge, [cond_nconsts, body_nconsts])\n  cond_consts, body_consts, carry = split_list(args, [cond_nconsts, body_nconsts])\n  cond_consts_avals, body_consts_avals, carry_avals = split_list(ctx.in_avals,\n                                                                 [cond_nconsts,\n                                                                  body_nconsts])\n\n  # Check if the same Ref is written to in both cond and body.\n  cond_write_ids = {id(cond_consts_avals[effect.input])\n    for effect in core.positional_effects(cond_jaxpr)\n    if isinstance(effect, state.WriteEffect)}\n  cond_has_writes = len(cond_write_ids) > 0\n  body_write_ids = {id(body_consts_avals[effect.input])\n    for effect in core.positional_effects(body_jaxpr)\n    if isinstance(effect, state.WriteEffect)}\n  write_to_both_ids = cond_write_ids & body_write_ids\n  if write_to_both_ids:\n    raise NotImplementedError(\n        \"Cannot write to the same ref in both cond and body of while loop.\")\n\n  cond_is_ref = [\n      isinstance(aval, state.AbstractRef) and should\n      for aval, should in zip(cond_consts_avals, cond_consts_discharge)\n  ]\n  remaining_cond_consts, cond_refs = partition_list(cond_is_ref, cond_consts)\n  remaining_cond_const_avals, cond_ref_avals = partition_list(cond_is_ref,\n                                                         cond_consts_avals)\n  num_cond_refs = sum(cond_is_ref)\n  num_remaining_cond_consts = cond_nconsts - num_cond_refs\n  body_is_ref = [\n      isinstance(aval, state.AbstractRef) and should\n      for aval, should in zip(body_consts_avals, body_consts_discharge)\n  ]\n  remaining_body_consts, body_refs = partition_list(body_is_ref, body_consts)\n  remaining_body_const_avals, body_ref_avals = partition_list(body_is_ref,\n                                                         body_consts_avals)","sourceCodeStart":2249,"sourceCodeEnd":2285,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2249-L2285","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","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"],"exampleFix":"// before\ndef cond(c):\n  ref.set(ref[...] + 1)\n  return c < 10\ndef body(c):\n  ref.set(ref[...] * 2)\n  return c + 1\n// after\ndef cond(c):\n  return c < 10\ndef body(c):\n  ref.set(ref[...] * 2)\n  return c + 1","handlingStrategy":"validation","validationCode":"# audit before compiling: ensure cond performs no ref writes\nimport jax\ndef cond_writes_refs(cond_fn, example_carry):\n    jaxpr = jax.make_jaxpr(cond_fn)(example_carry)\n    return any('Write' in type(x).__name__ or 'write' in str(x).lower() for x in jaxpr.effects)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","state-api","refs","while-loop","state-discharge"],"backgroundTag":"jax-ref-written-in-cond-and-body","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}