{"record":{"id":"3ac3e89bdf506c2e","repo":"jax-ml/jax","slug":"body-jaxpr-has-consts-if-you-see-this-error-plea","errorCode":null,"errorMessage":"Body jaxpr has consts. If you see this error, please open an issue at https://github.com/jax-ml/jax/issues","messagePattern":"Body jaxpr has consts\\. If you see this error, please open an issue at https://github\\.com/jax-ml/jax/issues","errorType":"error_code","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":2299,"sourceCode":"      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)\n  num_body_refs = sum(body_is_ref)\n  num_remaining_body_consts = body_nconsts - num_body_refs\n  num_out_body_consts = num_remaining_body_consts\n  if cond_has_writes:\n    # If the cond has writes, we need to add the cond consts into the body\n    # consts since we need to evaluate the cond condition in the body.\n    remaining_body_consts = [*remaining_cond_consts, *remaining_body_consts]\n    remaining_body_const_avals = [*remaining_cond_const_avals,\n                                  *remaining_body_const_avals]\n    num_remaining_body_consts += num_remaining_cond_consts\n\n  num_carry = len(ctx.in_avals) - body_nconsts - cond_nconsts\n  if body_jaxpr.consts:\n    raise NotImplementedError(\"Body jaxpr has consts. If you see this error, \"\n                              \"please open an issue at \"\n                              \"https://github.com/jax-ml/jax/issues\")\n  if cond_jaxpr.consts:\n    raise NotImplementedError(\"Cond jaxpr has consts. If you see this error, \"\n                              \"please open an issue at \"\n                              \"https://github.com/jax-ml/jax/issues\")\n  discharged_cond_jaxpr = state_discharge.discharge_state(\n      cond_jaxpr, should_discharge=[*cond_consts_discharge, *carry_discharge]\n  )\n  if discharged_cond_jaxpr.consts:\n    raise NotImplementedError\n  # body_jaxpr has the signature (*body_consts, *carry) -> carry.\n  # Some of these body_consts are actually `Ref`s so when we discharge\n  # them, they also turn into outputs, effectively turning those consts into\n  # carries. However this doesn't fit the expected signature for the body_jaxpr.\n  # Therefore we need to rewrite the jaxpr to shuffle around the `Ref`s so that\n  # they are part of the carry.\n  discharged_body_jaxpr = state_discharge.discharge_state(","sourceCodeStart":2281,"sourceCodeEnd":2317,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2281-L2317","documentation":"While discharging state from a while_loop, JAX requires the body jaxpr to have no closed-over constant cells remaining. If consts survive (e.g. refs that were not discharged or constants captured by tracing), discharge cannot proceed and JAX treats it as an internal/unhandled case, asking users to file an issue.","triggerScenarios":"Running the state-discharge transformation (io_effect discharge, jit lowering of stateful code) on a while_loop whose body jaxpr still has non-empty consts after discharge decisions — typically when some consts are refs partially discharged.","commonSituations":"Experimental stateful JAX (io_effect / state API) where the body closure captures refs or arrays; version mismatches between jax.experimental modules and core; rare internal path — most users hit it only via experimental features.","solutions":["Update JAX to the latest version — this is an internal invariant and may already be fixed","Minimize the repro (remove captured constants from the body function; pass everything explicitly as loop carry or body constants via closure-free code) and file an issue at https://github.com/jax-ml/jax/issues","Avoid io_effect/state discharge for this loop: rewrite using explicit carry instead of refs"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    f = jax.jit(stateful_while_fn)\n    out = f(args)\nexcept NotImplementedError as e:\n    if 'Body jaxpr has consts' in str(e):\n        # rewrite without refs / io_effect; use explicit carry\n        out = plain_carry_version(args)","preventionTips":["Avoid jax.experimental state/io_effect in production while_loops","Pass all values explicitly (carry/constants), avoid closure captures in body","Keep JAX updated; report internal-invariant errors upstream with a repro"],"tags":["jax","state-discharge","internal-error","while-loop"],"backgroundTag":"jax-internal-invariant-violation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}