jax-ml/jax · error · NotImplementedError

run_scoped_p with collective axes is not supported

Error message

run_scoped_p with collective axes is not supported

What it means

The Pallas TPU interpreter does not implement run_scoped_p with collective axes (collective run_scoped, where the scoped region spans multiple cores). Executing a mosaic run_scoped with a non-empty collective_axes parameter under interpret mode raises NotImplementedError.

Source

Thrown at jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py:1462

        def _body(val):
          token, val, _ = val
          token, val = _interpret(
              eqn.params['body_jaxpr'], *body_consts, *val, token=token)
          token, cond = _interpret(
              eqn.params['cond_jaxpr'], *cond_consts, *val, token=token)
          return token, val, cond[0]
        token, out, _ = lax.while_loop(
            lambda args: args[2], _body, (token, init_val, first_cond[0]))

      elif prim is pjit.jit_p:
        invals = deferred_invals()
        token, out = _interpret(eqn.params['jaxpr'],
                                *eqn.params['jaxpr'].consts,
                                *invals, token=token)

      elif prim is primitives.run_scoped_p:
        if eqn.params['collective_axes']:
          raise NotImplementedError(
              'run_scoped_p with collective axes is not supported'
          )
        # Allocate a buffer or semaphore for each element of
        # eqn.params['jaxpr'].invars. It is assumed that each core
        # runs the same sequence of `run_scoped`s.
        allocs = []
        for v in eqn.params['jaxpr'].invars:
          if v.aval.memory_space is _SEMAPHORE:
            token, alloc = callback.io_callback(
                _allocate_semaphores,
                (TOKEN_SHAPE_DTYPE,
                 jax.ShapeDtypeStruct(v.aval.shape, jnp.int16)),
                token,
                ctx.device_id,
                ctx.local_core_id,
                v.aval.shape,
            )
            allocs.append(alloc)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Run on real TPU hardware/compilation instead of interpret mode if collective run_scoped is essential
Refactor to per-core run_scoped without collective_axes (e.g., separate semaphores per core)
Check for interpreter updates in newer JAX versions that add collective support

Example fix

// not refactorable generically; use hardware execution
# before: interpret=True with collective run_scoped
# after: compile-and-run on TPU (interpret=False or actual device run)

Defensive patterns

Strategy: fallback

Validate before calling

# detect collective usage before interpret run
assert not kernel_uses_collective_run_scoped, 'collective run_scoped unsupported in interpret mode'

Try / catch

try:
    interpret_run(kernel)
except NotImplementedError as e:
    if 'collective axes' in str(e):
        run_on_tpu_hardware(kernel)

Prevention

Track which mosaic primitives the interpreter supports before adopting them
Gate interpret-mode debugging to non-collective code paths
Maintain a hardware execution path for collective features

When it happens

Trigger: Using mosaic_primitives.run_scoped(..., collective_axes=...) (collective semaphores / cross-core scoped regions) in a kernel interpreted with the TPU interpret mode.

Common situations: TPU kernels relying on cross-core collectives (e.g., collective DMA or cross-core semaphores) being debugged in interpret mode; newer mosaic features not yet ported to the interpreter.

Related errors

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