{"record":{"id":"25ba01100f8af477","repo":"jax-ml/jax","slug":"mesh-mesh-is-not-supported-by-the-mosaic-gpu-bac","errorCode":null,"errorMessage":"Mesh {mesh} is not supported by the Mosaic GPU backend","messagePattern":"Mesh (.+?) 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":63,"sourceCode":"    input_output_aliases: tuple[tuple[int, int], ...],\n    grid_mapping: pallas_core.GridMapping,\n    mesh: pallas_core.Mesh | None,\n    compiler_params: pallas_core.CompilerParams | None,\n    cost_estimate: pallas_core.CostEstimate | None,\n    out_avals: tuple[jax_core.AbstractValue, ...],\n    metadata: frozen_dict.FrozenDict[str, str] | None,\n    name: str | None,\n):\n  del metadata, name  # TODO(sharadmv): Add metadata to HLO.\n  debug_info = jaxpr.debug_info\n  del interpret, out_avals\n  if grid_mapping.num_dynamic_grid_bounds:\n    raise NotImplementedError(\n        \"dynamic grid bounds not supported in the Mosaic GPU backend\"\n    )\n\n  if mesh is not None and not isinstance(mesh, gpu_core.Mesh):\n    raise NotImplementedError(\n        f\"Mesh {mesh} is not supported by the Mosaic GPU backend\"\n    )\n\n  if debug:\n    print(f\"\\nThe kernel jaxpr for pallas_call {debug_info.func_src_info}:\")\n    print(jaxpr)\n    print(f\"The grid mapping for pallas_call {debug_info.func_src_info}:\")\n    print(grid_mapping)\n\n  mgpu.dialect.register_dialect(ctx.module_context.context)\n\n  if compiler_params is None:\n    gpu_params = gpu_core.CompilerParams()\n  else:\n    assert isinstance(compiler_params, gpu_core.CompilerParams)\n    gpu_params = compiler_params\n\n  jax_mesh = None","sourceCodeStart":45,"sourceCodeEnd":81,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py#L45-L81","documentation":"pallas_call on the Mosaic GPU backend only accepts meshes of type gpu_core.Mesh (or no mesh). Passing any other mesh object (e.g. a jax.sharding Mesh from a different backend or a custom object) raises NotImplementedError.","triggerScenarios":"Passing a mesh argument that is not None and not a gpu_core.Mesh instance to a pallas_call compiled for Mosaic GPU, e.g. a TPU-oriented Mesh or jax Mesh.","commonSituations":"Sharing kernel-launch helper code between TPU and GPU backends; accidentally passing jax.sharding.Mesh where plgpu.Mesh is expected.","solutions":["Pass None for mesh, or construct the correct gpu_core.Mesh","Branch per backend in shared launch code"],"exampleFix":"# before\nkernel[grid, mesh=jmesh](...)  # jax.sharding.Mesh\n# after\nkernel[grid](...)  # or pass gpu_core.Mesh instance","handlingStrategy":"validation","validationCode":"from jax._src.pallas import gpu_core\nassert mesh is None or isinstance(mesh, gpu_core.Mesh)","typeGuard":"def is_gpu_mesh(mesh) -> bool:\n    from jax._src.pallas import gpu_core\n    return mesh is None or isinstance(mesh, gpu_core.Mesh)","tryCatchPattern":null,"preventionTips":["Branch launch helpers per backend","Don't pass jax.sharding.Mesh to pallas GPU calls"],"tags":["pallas","mosaic-gpu","mesh","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}