{"record":{"id":"7ac22cf8aef548e8","repo":"jax-ml/jax","slug":"run-scoped-p-with-collective-axes-is-not-supported","errorCode":null,"errorMessage":"run_scoped_p with collective axes is not supported","messagePattern":"run_scoped_p with collective axes is not supported","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":1462,"sourceCode":"        def _body(val):\n          token, val, _ = val\n          token, val = _interpret(\n              eqn.params['body_jaxpr'], *body_consts, *val, token=token)\n          token, cond = _interpret(\n              eqn.params['cond_jaxpr'], *cond_consts, *val, token=token)\n          return token, val, cond[0]\n        token, out, _ = lax.while_loop(\n            lambda args: args[2], _body, (token, init_val, first_cond[0]))\n\n      elif prim is pjit.jit_p:\n        invals = deferred_invals()\n        token, out = _interpret(eqn.params['jaxpr'],\n                                *eqn.params['jaxpr'].consts,\n                                *invals, token=token)\n\n      elif prim is primitives.run_scoped_p:\n        if eqn.params['collective_axes']:\n          raise NotImplementedError(\n              'run_scoped_p with collective axes is not supported'\n          )\n        # Allocate a buffer or semaphore for each element of\n        # eqn.params['jaxpr'].invars. It is assumed that each core\n        # runs the same sequence of `run_scoped`s.\n        allocs = []\n        for v in eqn.params['jaxpr'].invars:\n          if v.aval.memory_space is _SEMAPHORE:\n            token, alloc = callback.io_callback(\n                _allocate_semaphores,\n                (TOKEN_SHAPE_DTYPE,\n                 jax.ShapeDtypeStruct(v.aval.shape, jnp.int16)),\n                token,\n                ctx.device_id,\n                ctx.local_core_id,\n                v.aval.shape,\n            )\n            allocs.append(alloc)","sourceCodeStart":1444,"sourceCodeEnd":1480,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L1444-L1480","documentation":"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.","triggerScenarios":"Using mosaic_primitives.run_scoped(..., collective_axes=...) (collective semaphores / cross-core scoped regions) in a kernel interpreted with the TPU interpret mode.","commonSituations":"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.","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"],"exampleFix":"// not refactorable generically; use hardware execution\n# before: interpret=True with collective run_scoped\n# after: compile-and-run on TPU (interpret=False or actual device run)","handlingStrategy":"fallback","validationCode":"# detect collective usage before interpret run\nassert not kernel_uses_collective_run_scoped, 'collective run_scoped unsupported in interpret mode'","typeGuard":null,"tryCatchPattern":"try:\n    interpret_run(kernel)\nexcept NotImplementedError as e:\n    if 'collective axes' in str(e):\n        run_on_tpu_hardware(kernel)","preventionTips":["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"],"tags":["jax","pallas","tpu","interpret-mode","not-implemented","collectives"],"backgroundTag":"unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}