{"record":{"id":"cdc55b4b9039cd82","repo":"jax-ml/jax","slug":"state-effect-not-supported-in-cond-partial-eval","errorCode":null,"errorMessage":"State effect not supported in cond partial-eval.","messagePattern":"State effect not supported in cond partial-eval\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/conditionals.py","lineNumber":609,"sourceCode":"  index, *ops = primals\n  _, *ops_dot = tangents\n  ops_dot = _prune_zeros(ops_dot)\n\n  out = cond_p.bind(index, *ops, *ops_dot, branches=branches_jvp,\n                    **params)\n  out_primals, out_tangents = split_list(out, [len(out_nz)])\n  out_tangents_iter = iter(out_tangents)\n  out_tangents = [next(out_tangents_iter) if nz else\n                  ad_util.p2tz(p)\n                  for p, nz in zip(out_primals, out_nz)]\n  return out_primals, out_tangents\n\ndef _cond_partial_eval(trace, *tracers, branches, **params):\n  in_unknowns = [not t.pval.is_known() for t in tracers]\n  index_uk, *ops_uk = in_unknowns\n  if any(isinstance(eff, RefEffect) for branch in branches for eff in\n      branch.effects):\n    raise NotImplementedError(\n        \"State effect not supported in cond partial-eval.\")\n\n  if index_uk:\n    # When the branch index is unknown, we stage out the whole cond.\n    # TODO(mattjj): remove this path when old remat is removed\n    params = dict(branches=branches, **params)\n    return trace.default_process_primitive(cond_p, tracers, params)\n\n  branches_out_uks = []\n  for branch_jaxpr in branches:\n    _, _, out_uks, _ = pe.partial_eval_jaxpr_nounits(\n        branch_jaxpr, ops_uk, instantiate=False)\n    branches_out_uks.append(out_uks)\n  out_uks = [any(uks) for uks in zip(*branches_out_uks)]\n\n  branches_known, branches_unknown, branch_res_avals = [], [], []\n  for branch_jaxpr in branches:\n    branch_jaxpr_known, branch_jaxpr_unknown, _, res_avals = \\","sourceCodeStart":591,"sourceCodeEnd":627,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/conditionals.py#L591-L627","documentation":"Raised during partial evaluation of lax.cond when a branch contains RefEffects (state mutation). Partial eval cannot stage half of a stateful cond, because unknown values would leave the ref updates undefined, so it errors.","triggerScenarios":"Running gradient (grad), jit with unknown-constant inputs, or any transform that triggers pe.trace_to_jaxpr over a cond whose branches mutate refs.","commonSituations":"Differentiating or rematerializing through experimental stateful code that uses lax.cond; combining new-style state with jax.grad or custom JVPs.","solutions":["Make branches pure: read refs before the cond and write after it, selecting values with jnp.where","Avoid autodiff over state-mutating conds; restructure to update state outside conditional logic","Upgrade JAX: newer versions progressively support state effects in more transforms","Fall back to classic lax.fori_loop/scan with explicit carry instead of refs"],"exampleFix":"// before\nval = ref_get(ref)\nout = lax.cond(flag, lambda: val + 1, lambda: val)\nref_set(ref, out)\n# inside grad-traced code with refs in cond branches\n// after\nval = ref_get(ref)\nout = jnp.where(flag, val + 1, val)  # autodiff-safe\nref_set(ref, out)","handlingStrategy":"fallback","validationCode":"null","typeGuard":"null","tryCatchPattern":"try: jax.grad(f)(x)\\nexcept NotImplementedError as e:\\n    if 'State effect not supported in cond partial-eval' in str(e):\\n        return functional_version(f)(x)  # ref-free reimplementation\\n    raise","preventionTips":["Keep refs out of cond branches whenever autodiff is involved","Maintain a pure functional mirror of stateful experimental code","Check effect support matrix before combining state with transforms"],"tags":["jax","cond","partial-eval","state-effects","autodiff"],"backgroundTag":"jax-unsupported-effect-in-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}