{"record":{"id":"d5a95ad9012ef145","repo":"jax-ml/jax","slug":"run-scoped-discharge-does-not-support-collective-a","errorCode":null,"errorMessage":"run_scoped discharge does not support collective_axes yet.","messagePattern":"run_scoped discharge does not support collective_axes yet\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/primitives.py","lineNumber":755,"sourceCode":"  for eff in jaxpr.effects:\n    if isinstance(eff, effects.JaxprInputEffect):\n      if eff.input in constvar_idx:\n        nonlocal_effects.add(eff.replace(constvar_idx[eff.input]))\n      continue\n    nonlocal_effects.add(eff)\n  return [v.aval for v in jaxpr.outvars], nonlocal_effects\n\n\ndef _run_scoped_discharge_rule(\n    ctx,\n    *args_flat,\n    jaxpr,\n    collective_axes,\n    ref_transforms,\n    **_,\n):\n  if collective_axes:\n    raise NotImplementedError(\n        \"run_scoped discharge does not support collective_axes yet.\"\n    )\n  # discharge_state only discharges invars, not consts, so in order to\n  # discharge the requested refs we need to move them to the invar set.\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,\n      should_discharge=ctx.should_discharge + [False] * len(jaxpr.invars),\n  )\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  # Lowering expects that the jaxpr.consts to be the eqn.invals.\n  discharged_body = discharged_body.with_consts(args_flat)\n","sourceCodeStart":737,"sourceCodeEnd":773,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/primitives.py#L737-L773","documentation":"Pallas run_scoped cannot discharge (lift out of scoped execution) a function that uses collective axes (multi-device/multi-core collectives). State discharge for collective operations inside run_scoped is unimplemented in JAX Pallas.","triggerScenarios":"Passing a non-empty collective_axes argument to run_scoped which then goes through the state-discharge path (e.g., interpret mode or non-Pallas lowering).","commonSituations":"Porting distributed/multi-device TPU kernel code that uses collectives into a run_scoped block; running a kernel with collectives outside an actual Pallas call context where discharge is required.","solutions":["Remove collective_axes from the run_scoped call if collectives are not needed","Keep collective operations outside the run_scoped body and pass results in as arguments","File or star the upstream JAX issue requesting collective discharge support and restructure until then"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert not collective_axes, \"collective_axes unsupported by run_scoped discharge\"","typeGuard":null,"tryCatchPattern":"try:\n    run_scoped(f, refs, collective_axes=axes)\nexcept NotImplementedError:\n    # fall back to non-collective version\n    run_scoped(f, refs)","preventionTips":["Check collective_axes is empty before run_scoped in interpret/CPU paths","Keep collectives outside run_scoped bodies"],"tags":["pallas","run-scoped","collectives","not-implemented","jax"],"backgroundTag":"unsupported-feature-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}