{"record":{"id":"c02e90c747a01052","repo":"jax-ml/jax","slug":"cond-jaxpr-has-consts-if-you-see-this-error-plea","errorCode":null,"errorMessage":"Cond jaxpr has consts. If you see this error, please open an issue at https://github.com/jax-ml/jax/issues","messagePattern":"Cond 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":2303,"sourceCode":"                                                         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(\n      body_jaxpr, should_discharge=[*body_consts_discharge, *carry_discharge]\n  )\n  if discharged_body_jaxpr.consts:\n    raise NotImplementedError","sourceCodeStart":2285,"sourceCodeEnd":2321,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L2285-L2321","documentation":"Companion check to the body-consts one: while discharging state from a while_loop, the cond jaxpr must have no remaining consts. If constants survive discharge, the transformation cannot proceed and JAX raises NotImplementedError requesting an issue be filed, since this indicates an unhandled internal case.","triggerScenarios":"Running state discharge on a while_loop whose cond jaxpr still has consts after discharge — e.g. the cond closure captures arrays or partially-discharged refs.","commonSituations":"Experimental stateful JAX code where the cond function closes over refs/values; typically only reachable through jax.experimental state/io_effect features or Pallas pipelines; may indicate a JAX version bug.","solutions":["Update JAX to the latest version (internal invariant, likely fixed upstream)","Rewrite the cond to take all values as explicit loop carry or cond constants rather than Python closure captures","File an issue at https://github.com/jax-ml/jax/issues with a minimal reproduction"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    out = jax.jit(stateful_while_fn)(args)\nexcept NotImplementedError as e:\n    if 'Cond jaxpr has consts' in str(e):\n        out = plain_carry_version(args)  # rewrite cond closure-free","preventionTips":["Write cond functions that reference only their arguments","Avoid mixing experimental ref state with while_loop cond","Update JAX before debugging further; file an issue with a minimal 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"}