jax-ml/jax · error · ValueError

run_scoped lowering outside of Pallas does not support colle

Error message

run_scoped lowering outside of Pallas does not support collective_axes.

What it means

run_scoped lowering outside an actual Pallas kernel execution does not support collective_axes. Using run_scoped with collective axes in a non-Pallas context (e.g., plain jit) is a ValueError.

Source

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

  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 = [
      ref_outputs.pop(0) if should and isinstance(aval, state.AbstractRef)
      else None for should, aval in zip(ctx.should_discharge, ctx.in_avals)]
  assert len(updates) == len(ctx.in_avals), f'{len(updates)} != {len(ctx.in_avals)}'
  return updates, return_values


state_discharge.register_discharge_rule(run_scoped_p)(
    _run_scoped_discharge_rule)


@functools.partial(mlir.register_lowering, run_scoped_p)
def _run_scoped_lowering_rule(ctx, *args, jaxpr, collective_axes, **_):
  if collective_axes:
    raise ValueError(
        "run_scoped lowering outside of Pallas does not support"
        " collective_axes."
    )
  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=True)
  discharged_body, new_consts = discharged_closed_body, discharged_closed_body.consts
  if new_consts:
    raise NotImplementedError(
        "Cannot handle new consts created by state discharge.")

  def _lower_fun(*lower_fun_args):
    num_consts = len(lower_fun_args)
    body_avals = [v.aval for v in discharged_body.invars[num_consts:]]
    # Create inputs filled with uninitialized values to the body.
    init_vals = [
        uninitialized_value(aval.shape, aval.dtype) for aval in body_avals  # type: ignore

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove collective_axes when calling run_scoped outside a Pallas kernel
Run the code inside a proper Pallas call (pallas_call / pallas kernel) where collectives are supported
Conditionally pass collective_axes only when executing on TPU with Pallas

Example fix

// before
run_scoped(f, refs, collective_axes="i")
// after (outside Pallas)
run_scoped(f, refs)

Defensive patterns

Strategy: validation

Validate before calling

in_pallas = is_inside_pallas_context()  # your tracking flag
if collective_axes and not in_pallas:
    collective_axes = None  # or raise early with a clear message

Try / catch

try:
    run_scoped(f, refs, collective_axes=axes)
except ValueError as e:
    if "collective_axes" in str(e):
        axes = None
        run_scoped(f, refs)
    raise

Prevention

Only pass collective_axes inside real Pallas execution
Gate distributed code paths on backend type

When it happens

Trigger: Calling run_scoped(f, refs, collective_axes=...) under plain jax.jit or any non-Pallas lowering context.

Common situations: Testing Pallas kernel helper code under plain jit; migrating code from pallas_call/TPU kernel context to CPU/interpret mode while keeping collective_axes.

Related errors

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