{"record":{"id":"14063542cad4ba04","repo":"jax-ml/jax","slug":"lowering-multiple-mesh-function-pairs-is-not-suppo","errorCode":null,"errorMessage":"Lowering multiple mesh/function pairs is not supported by the Mosaic GPU backend","messagePattern":"Lowering multiple mesh/function pairs is not supported by the Mosaic GPU backend","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/pallas_call_registration.py","lineNumber":199,"sourceCode":"def mpmd_map_mgpu_lowering_rule(\n    ctx: mlir.LoweringRuleContext,\n    *args,\n    meshes,\n    jaxprs,\n    out_avals,\n    input_output_aliases,\n    compiler_params,\n    interpret,\n    debug,\n    cost_estimate,\n    metadata,\n    name,\n    external_meshes,\n):\n  del interpret, cost_estimate, metadata, name, out_avals  # Unused.\n\n  if len(jaxprs) != 1:\n    raise NotImplementedError(\n        \"Lowering multiple mesh/function pairs is not supported by the Mosaic\"\n        \" GPU backend\"\n    )\n  if external_meshes:\n    raise NotImplementedError(\n        \"External meshes are not supported by the Mosaic GPU backend\"\n    )\n  [jaxpr] = jaxprs\n  [mesh] = meshes\n  if not isinstance(mesh, gpu_core.Mesh):\n    raise NotImplementedError(\n        f\"Mesh {mesh} is not supported by the Mosaic GPU backend\"\n    )\n  # On GPU ``mpmd_map`` kernels never carry scratch operands -- scratch is\n  # handled separately by ``plgpu.kernel``. So the jaxpr invars are exactly the\n  # inputs followed by the outputs.\n  if len(jaxpr.invars) != len(args) + len(ctx.avals_out):\n    raise NotImplementedError(","sourceCodeStart":181,"sourceCodeEnd":217,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py#L181-L217","documentation":"The mpmd_map lowering for Mosaic GPU only supports exactly one (jaxpr, mesh) function pair. If mpmd_map is given multiple functions to lower in an MPMD fashion, this NotImplementedError is raised since multi-program multi-data lowering is unimplemented.","triggerScenarios":"Calling jax.mpmd_map (or the mpmd_map primitive) with a list of more than one function/mesh pair while targeting the Mosaic GPU backend.","commonSituations":"Reusing MPMD pipelines written for TPU on GPU; building multi-kernel launch configurations.","solutions":["Split the mpmd_map into separate single-function kernel launches","Await or contribute upstream support for multi-pair mpmd_map on GPU"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":"assert len(jaxprs) == 1, \"mpmd_map on GPU supports exactly one function pair\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Launch one kernel per function on GPU"],"tags":["pallas","mosaic-gpu","mpmd-map","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}