jax-ml/jax · error · NotImplementedError

run_scoped discharge does not support collective_axes yet.

Error message

run_scoped discharge does not support collective_axes yet.

What it means

Pallas run_scoped cannot discharge (lift out of scoped execution) a function that uses collective axes (multi-device/multi-core collectives). State discharge for collective operations inside run_scoped is unimplemented in JAX Pallas.

Source

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

  for eff in jaxpr.effects:
    if isinstance(eff, effects.JaxprInputEffect):
      if eff.input in constvar_idx:
        nonlocal_effects.add(eff.replace(constvar_idx[eff.input]))
      continue
    nonlocal_effects.add(eff)
  return [v.aval for v in jaxpr.outvars], nonlocal_effects


def _run_scoped_discharge_rule(
    ctx,
    *args_flat,
    jaxpr,
    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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove collective_axes from the run_scoped call if collectives are not needed
2. Keep collective operations outside the run_scoped body and pass results in as arguments
3. File or star the upstream JAX issue requesting collective discharge support and restructure until then

assert not collective_axes, "collective_axes unsupported by run_scoped discharge"

try:
    run_scoped(f, refs, collective_axes=axes)
except NotImplementedError:
    # fall back to non-collective version
    run_scoped(f, refs)

When it happens

Trigger: Passing a non-empty collective_axes argument to run_scoped which then goes through the state-discharge path (e.g., interpret mode or non-Pallas lowering).

Common situations: Porting distributed/multi-device TPU kernel code that uses collectives into a run_scoped block; running a kernel with collectives outside an actual Pallas call context where discharge is required.

Related errors

• run_scoped_p with collective axes is not supported
• Cannot handle new consts created by state discharge.
• run_scoped lowering outside of Pallas does not support colle
• {axis_name} mixes JAX mesh and Pallas mesh grid axes
• Explicit sharding is not currently supported in the pallas-t

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