{"record":{"id":"9a48a870fe315b63","repo":"jax-ml/jax","slug":"run-scoped-interpret-rule-does-not-support-collect","errorCode":null,"errorMessage":"run_scoped interpret rule does not support collective axes","messagePattern":"run_scoped interpret rule does not support collective axes","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/hlo_interpreter.py","lineNumber":287,"sourceCode":"        mapped_jaxprs, mapped_args = zip(*map(\n          lambda x, i: _resolve_jaxpr(interpreter, x, mapped_idx=i), value, range(len(value))))\n        all_new_args = tuple(new_arg for _args in mapped_args for new_arg in _args)\n        new_params[key] = tuple(mapped_jaxprs)\n        args = all_new_args + args\n      else:\n        raise ValueError(f\"Parameter {key} is not a Jaxpr or sequence of Jaxprs: {value}\")\n    params.update(new_params)\n    return primitive.bind(*args, **params)\n  return rule\n\n_eval_jaxpr_hop_rules[loops.scan_p] = make_hop_rule(loops.scan_p, 'jaxpr')\n_eval_jaxpr_hop_rules[loops.while_p] = make_hop_rule(\n    loops.while_p, 'body_jaxpr', 'cond_jaxpr')\n_eval_jaxpr_hop_rules[conditionals.cond_p] = make_hop_rule(conditionals.cond_p, 'branches')\ndef _run_scoped_physicalize_rule(\n    interpreter, *consts, jaxpr: jax_core.Jaxpr, collective_axes, **params):\n  if collective_axes:\n    raise NotImplementedError(\n        \"run_scoped interpret rule does not support collective axes\"\n    )\n  physical_jaxpr, physical_consts = interpreter(jaxpr, consts)\n  return primitives.run_scoped_p.bind(\n      *physical_consts, jaxpr=physical_jaxpr, collective_axes=collective_axes,\n      **params\n  )\n_eval_jaxpr_hop_rules[primitives.run_scoped_p] = _run_scoped_physicalize_rule\n\n\n# TODO(justinfu): Replace this with a standardized physicalize pass.\ndef resolve_physical_types(jaxpr: jax_core.Jaxpr, consts: Sequence[Any]):\n  kernel_avals = jaxpr.in_avals\n  kernel_avals = tuple(map(_logical_aval_to_interpret_mode_aval,\n                             kernel_avals))\n  interp_fun = partial(\n      eval_jaxpr_recursive, jaxpr, consts,\n      recurse_hop_rule=resolve_physical_types)","sourceCodeStart":269,"sourceCodeEnd":305,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/hlo_interpreter.py#L269-L305","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Drop the collective axes (remove the mesh/collective ops from the scoped region) before interpreting","Test with the real compiler/backend instead of the HLO interpreter when collectives are required","Restructure so collective operations happen outside the run_scoped region interpreted by pallas"],"exampleFix":"// before\nwith jax.lax.run_scoped(..., collective_axes=('mesh',)):\n  ...collective ops...\n// after\nwith jax.lax.run_scoped(..., collective_axes=()):\n  ...local ops only...","handlingStrategy":"type-guard","validationCode":"def has_no_collective_axes(f):\n    # inspect traced jaxpr before interpret\n    jaxpr = jax.make_jaxpr(f)()\n    return all(eqn.params.get('collective_axes', ()) == () for eqn in jaxpr.eqns)","typeGuard":null,"tryCatchPattern":"try:\n    interp_run(kernel)\nexcept NotImplementedError as e:\n    if 'collective axes' in str(e):\n        raise RuntimeError('Use compiled mode for collective kernels') from e\n    raise","preventionTips":["Grep kernels for run_scoped with collective_axes before interpret runs","Keep an interpreter-compatible and a compiled variant of mesh kernels"],"tags":["jax","pallas","collectives","interpreter","notimplemented"],"backgroundTag":"unsupported-feature-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}