{"record":{"id":"1f161f706e17b2b5","repo":"jax-ml/jax","slug":"external-meshes-are-not-supported-by-the-mosaic-gp","errorCode":null,"errorMessage":"External meshes are not supported by the Mosaic GPU backend","messagePattern":"External meshes are 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":204,"sourceCode":"    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(\n        \"Scratch operands are not supported by the Mosaic GPU mpmd_map lowering\"\n    )\n\n  if debug:\n    print(f\"\\nThe kernel jaxpr for mpmd_map {jaxpr.debug_info.func_src_info}:\")","sourceCodeStart":186,"sourceCodeEnd":222,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py#L186-L222","documentation":"mpmd_map on Mosaic GPU does not support external meshes (meshes passed via the external_meshes mechanism). Only the mesh bound directly with the jaxpr is accepted.","triggerScenarios":"Invoking mpmd_map with external_meshes non-empty while compiling for Mosaic GPU.","commonSituations":"Advanced multi-device setups ported from other backends that rely on externally supplied mesh resources.","solutions":["Remove external_meshes and bind the mesh directly to the computation","Restructure to a single-kernel launch with its own gpu_core.Mesh"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert not external_meshes","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Bind meshes directly; avoid the external_meshes path on GPU"],"tags":["pallas","mosaic-gpu","mpmd-map","mesh","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}