jax-ml/jax · error · NotImplementedError

Cond jaxpr has consts. If you see this error, please open an

Error message

Cond jaxpr has consts. If you see this error, please open an issue at https://github.com/jax-ml/jax/issues

What it means

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.

Source

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

                                                         body_consts_avals)
  num_body_refs = sum(body_is_ref)
  num_remaining_body_consts = body_nconsts - num_body_refs
  num_out_body_consts = num_remaining_body_consts
  if cond_has_writes:
    # If the cond has writes, we need to add the cond consts into the body
    # consts since we need to evaluate the cond condition in the body.
    remaining_body_consts = [*remaining_cond_consts, *remaining_body_consts]
    remaining_body_const_avals = [*remaining_cond_const_avals,
                                  *remaining_body_const_avals]
    num_remaining_body_consts += num_remaining_cond_consts

  num_carry = len(ctx.in_avals) - body_nconsts - cond_nconsts
  if body_jaxpr.consts:
    raise NotImplementedError("Body jaxpr has consts. If you see this error, "
                              "please open an issue at "
                              "https://github.com/jax-ml/jax/issues")
  if cond_jaxpr.consts:
    raise NotImplementedError("Cond jaxpr has consts. If you see this error, "
                              "please open an issue at "
                              "https://github.com/jax-ml/jax/issues")
  discharged_cond_jaxpr = state_discharge.discharge_state(
      cond_jaxpr, should_discharge=[*cond_consts_discharge, *carry_discharge]
  )
  if discharged_cond_jaxpr.consts:
    raise NotImplementedError
  # body_jaxpr has the signature (*body_consts, *carry) -> carry.
  # Some of these body_consts are actually `Ref`s so when we discharge
  # them, they also turn into outputs, effectively turning those consts into
  # carries. However this doesn't fit the expected signature for the body_jaxpr.
  # Therefore we need to rewrite the jaxpr to shuffle around the `Ref`s so that
  # they are part of the carry.
  discharged_body_jaxpr = state_discharge.discharge_state(
      body_jaxpr, should_discharge=[*body_consts_discharge, *carry_discharge]
  )
  if discharged_body_jaxpr.consts:
    raise NotImplementedError

View on GitHub (pinned to 1e1c6a8fc0)

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

Defensive patterns

Strategy: fallback

Try / catch

try:
    out = jax.jit(stateful_while_fn)(args)
except NotImplementedError as e:
    if 'Cond jaxpr has consts' in str(e):
        out = plain_carry_version(args)  # rewrite cond closure-free

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Related errors

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