jax-ml/jax · error · NotImplementedError

run_scoped interpret rule does not support collective axes

Error message

run_scoped interpret rule does not support collective axes

What it means

The Pallas HLO interpreter's rule for run_scoped explicitly rejects computations declared with collective (multi-device/multi-core) axes, because the interpreter only simulates single-device execution. When a run_scoped primitive carries non-empty collective_axes, there is no simulated communication, so it raises NotImplementedError.

Source

Thrown at jax/_src/pallas/hlo_interpreter.py:287

        mapped_jaxprs, mapped_args = zip(*map(
          lambda x, i: _resolve_jaxpr(interpreter, x, mapped_idx=i), value, range(len(value))))
        all_new_args = tuple(new_arg for _args in mapped_args for new_arg in _args)
        new_params[key] = tuple(mapped_jaxprs)
        args = all_new_args + args
      else:
        raise ValueError(f"Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value}")
    params.update(new_params)
    return primitive.bind(*args, **params)
  return rule

_eval_jaxpr_hop_rules[loops.scan_p] = make_hop_rule(loops.scan_p, 'jaxpr')
_eval_jaxpr_hop_rules[loops.while_p] = make_hop_rule(
    loops.while_p, 'body_jaxpr', 'cond_jaxpr')
_eval_jaxpr_hop_rules[conditionals.cond_p] = make_hop_rule(conditionals.cond_p, 'branches')
def _run_scoped_physicalize_rule(
    interpreter, *consts, jaxpr: jax_core.Jaxpr, collective_axes, **params):
  if collective_axes:
    raise NotImplementedError(
        "run_scoped interpret rule does not support collective axes"
    )
  physical_jaxpr, physical_consts = interpreter(jaxpr, consts)
  return primitives.run_scoped_p.bind(
      *physical_consts, jaxpr=physical_jaxpr, collective_axes=collective_axes,
      **params
  )
_eval_jaxpr_hop_rules[primitives.run_scoped_p] = _run_scoped_physicalize_rule


# TODO(justinfu): Replace this with a standardized physicalize pass.
def resolve_physical_types(jaxpr: jax_core.Jaxpr, consts: Sequence[Any]):
  kernel_avals = jaxpr.in_avals
  kernel_avals = tuple(map(_logical_aval_to_interpret_mode_aval,
                             kernel_avals))
  interp_fun = partial(
      eval_jaxpr_recursive, jaxpr, consts,
      recurse_hop_rule=resolve_physical_types)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Drop the collective axes (remove the mesh/collective ops from the scoped region) before interpreting
  2. Test with the real compiler/backend instead of the HLO interpreter when collectives are required
  3. Restructure so collective operations happen outside the run_scoped region interpreted by pallas

Example fix

// before
with jax.lax.run_scoped(..., collective_axes=('mesh',)):
  ...collective ops...
// after
with jax.lax.run_scoped(..., collective_axes=()):
  ...local ops only...
Defensive patterns

Strategy: type-guard

Validate before calling

def has_no_collective_axes(f):
    # inspect traced jaxpr before interpret
    jaxpr = jax.make_jaxpr(f)()
    return all(eqn.params.get('collective_axes', ()) == () for eqn in jaxpr.eqns)

Try / catch

try:
    interp_run(kernel)
except NotImplementedError as e:
    if 'collective axes' in str(e):
        raise RuntimeError('Use compiled mode for collective kernels') from e
    raise

Prevention

When it happens

Trigger: Using the HLO interpreter on a kernel whose body was traced inside jax.lax.run_scoped with collective axes (e.g., mesh/collective operations inside a scoped region), or interpreting lowered TPU/Mosaic pipelines that contain run_scoped_p with collective_axes set.

Common situations: Running interpret/debug mode on SPMD or mesh-scoped Pallas kernels; lowering pipelines produced under jax.sharding.MapAxisResources that insert run_scoped with collectives; expecting the interpreter to emulate collectives like real compilation does.

Related errors

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