{"record":{"id":"db77717120352c31","repo":"jax-ml/jax","slug":"run-scoped-lowering-outside-of-pallas-does-not-sup","errorCode":null,"errorMessage":"run_scoped lowering outside of Pallas does not support collective_axes.","messagePattern":"run_scoped lowering outside of Pallas does not support collective_axes\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/primitives.py","lineNumber":800,"sourceCode":"  return_values = out[:num_return_values]\n  ref_outputs = out[num_return_values:]\n  # We update all ref values with their updated values from the discharged\n  # body. For other values we leave them in place.\n  updates = [\n      ref_outputs.pop(0) if should and isinstance(aval, state.AbstractRef)\n      else None for should, aval in zip(ctx.should_discharge, ctx.in_avals)]\n  assert len(updates) == len(ctx.in_avals), f'{len(updates)} != {len(ctx.in_avals)}'\n  return updates, return_values\n\n\nstate_discharge.register_discharge_rule(run_scoped_p)(\n    _run_scoped_discharge_rule)\n\n\n@functools.partial(mlir.register_lowering, run_scoped_p)\ndef _run_scoped_lowering_rule(ctx, *args, jaxpr, collective_axes, **_):\n  if collective_axes:\n    raise ValueError(\n        \"run_scoped lowering outside of Pallas does not support\"\n        \" collective_axes.\"\n    )\n  jaxpr_noconst = pe.convert_constvars_jaxpr(jaxpr)\n  num_return_values = len(jaxpr_noconst.outvars)\n  discharged_closed_body = state_discharge.discharge_state(\n      jaxpr_noconst, should_discharge=True)\n  discharged_body, new_consts = discharged_closed_body, discharged_closed_body.consts\n  if new_consts:\n    raise NotImplementedError(\n        \"Cannot handle new consts created by state discharge.\")\n\n  def _lower_fun(*lower_fun_args):\n    num_consts = len(lower_fun_args)\n    body_avals = [v.aval for v in discharged_body.invars[num_consts:]]\n    # Create inputs filled with uninitialized values to the body.\n    init_vals = [\n        uninitialized_value(aval.shape, aval.dtype) for aval in body_avals  # type: ignore","sourceCodeStart":782,"sourceCodeEnd":818,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/primitives.py#L782-L818","documentation":"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.","triggerScenarios":"Calling run_scoped(f, refs, collective_axes=...) under plain jax.jit or any non-Pallas lowering context.","commonSituations":"Testing Pallas kernel helper code under plain jit; migrating code from pallas_call/TPU kernel context to CPU/interpret mode while keeping collective_axes.","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"],"exampleFix":"// before\nrun_scoped(f, refs, collective_axes=\"i\")\n// after (outside Pallas)\nrun_scoped(f, refs)","handlingStrategy":"validation","validationCode":"in_pallas = is_inside_pallas_context()  # your tracking flag\nif collective_axes and not in_pallas:\n    collective_axes = None  # or raise early with a clear message","typeGuard":null,"tryCatchPattern":"try:\n    run_scoped(f, refs, collective_axes=axes)\nexcept ValueError as e:\n    if \"collective_axes\" in str(e):\n        axes = None\n        run_scoped(f, refs)\n    raise","preventionTips":["Only pass collective_axes inside real Pallas execution","Gate distributed code paths on backend type"],"tags":["pallas","run-scoped","collectives","jax"],"backgroundTag":"invalid-argument-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}