jax-ml/jax · error · NotImplementedError

Cannot handle new consts created by state discharge.

Error message

Cannot handle new consts created by state discharge.

What it means

When discharging state for run_scoped, the state_discharge transformation unexpectedly produced new constants (captured values) in the discharged jaxpr, which this lowering path cannot handle. Only pre-existing invars/consts are supported.

Source

Thrown at jax/_src/pallas/primitives.py:768

    collective_axes,
    ref_transforms,
    **_,
):
  if collective_axes:
    raise NotImplementedError(
        "run_scoped discharge does not support collective_axes yet."
    )
  # discharge_state only discharges invars, not consts, so in order to
  # discharge the requested refs we need to move them to the invar set.
  jaxpr_noconst = pe.convert_constvars_jaxpr(jaxpr)
  num_return_values = len(jaxpr_noconst.outvars)
  discharged_closed_body = state_discharge.discharge_state(
      jaxpr_noconst,
      should_discharge=ctx.should_discharge + [False] * len(jaxpr.invars),
  )
  discharged_body, new_consts = discharged_closed_body, discharged_closed_body.consts
  if new_consts:
    raise NotImplementedError(
        "Cannot handle new consts created by state discharge.")

  # Lowering expects that the jaxpr.consts to be the eqn.invals.
  discharged_body = discharged_body.with_consts(args_flat)

  # Run_scoped discharged the external variables but the scoped ones
  # are not discharged.
  out = run_scoped_p.bind(
      *args_flat, jaxpr=discharged_body, collective_axes=collective_axes,
      ref_transforms=ref_transforms,
  )
  # Order of outputs:
  # (1) return values, (2) closed refs, (3) scoped refs.
  return_values = out[:num_return_values]
  ref_outputs = out[num_return_values:]
  # We update all ref values with their updated values from the discharged
  # body. For other values we leave them in place.
  updates = [

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Simplify the run_scoped body so discharge does not create new constants (hoist computations out of run_scoped)
Update JAX to the latest version — this is an internal discharge limitation that gets fixed over time
Report the reproducer to the JAX team if it persists on latest

Defensive patterns

Strategy: fallback

Try / catch

try:
    out = run_scoped(f, refs)
except NotImplementedError as e:
    if "new consts" in str(e):
        # simplify body / hoist constants and retry
        out = run_scoped(simplified_f, refs)
    raise

Prevention

Keep run_scoped bodies simple; hoist constants out
Track JAX release notes for state-discharge fixes

When it happens

Trigger: A run_scoped body whose discharged jaxpr closes over new constants after discharge_state runs (e.g., derived constants created during discharge of scoped refs).

Common situations: Complex kernels where discharge materializes new constant values; typically a JAX version-specific internal limitation rather than a user API mistake.

Related errors

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