jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

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.

Source

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

      for aval, should in zip(body_consts_avals, body_consts_discharge)
  ]
  remaining_body_consts, body_refs = partition_list(body_is_ref, body_consts)
  remaining_body_const_avals, body_ref_avals = partition_list(body_is_ref,
                                                         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(

View on GitHub (pinned to 1e1c6a8fc0)

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

Defensive patterns

Strategy: fallback

Try / catch

try:
    f = jax.jit(stateful_while_fn)
    out = f(args)
except NotImplementedError as e:
    if 'Body jaxpr has consts' in str(e):
        # rewrite without refs / io_effect; use explicit carry
        out = plain_carry_version(args)

Prevention

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

When it happens

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

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

Related errors

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