jax-ml/jax · error · NotImplementedError

External meshes are not supported by the Mosaic GPU backend

Error message

External meshes are not supported by the Mosaic GPU backend

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/pallas_call_registration.py:204

    out_avals,
    input_output_aliases,
    compiler_params,
    interpret,
    debug,
    cost_estimate,
    metadata,
    name,
    external_meshes,
):
  del interpret, cost_estimate, metadata, name, out_avals  # Unused.

  if len(jaxprs) != 1:
    raise NotImplementedError(
        "Lowering multiple mesh/function pairs is not supported by the Mosaic"
        " GPU backend"
    )
  if external_meshes:
    raise NotImplementedError(
        "External meshes are not supported by the Mosaic GPU backend"
    )
  [jaxpr] = jaxprs
  [mesh] = meshes
  if not isinstance(mesh, gpu_core.Mesh):
    raise NotImplementedError(
        f"Mesh {mesh} is not supported by the Mosaic GPU backend"
    )
  # On GPU ``mpmd_map`` kernels never carry scratch operands -- scratch is
  # handled separately by ``plgpu.kernel``. So the jaxpr invars are exactly the
  # inputs followed by the outputs.
  if len(jaxpr.invars) != len(args) + len(ctx.avals_out):
    raise NotImplementedError(
        "Scratch operands are not supported by the Mosaic GPU mpmd_map lowering"
    )

  if debug:
    print(f"\nThe kernel jaxpr for mpmd_map {jaxpr.debug_info.func_src_info}:")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove external_meshes and bind the mesh directly to the computation
Restructure to a single-kernel launch with its own gpu_core.Mesh

Defensive patterns

Strategy: validation

Validate before calling

assert not external_meshes

Prevention

Bind meshes directly; avoid the external_meshes path on GPU

When it happens

Trigger: Invoking mpmd_map with external_meshes non-empty while compiling for Mosaic GPU.

Common situations: Advanced multi-device setups ported from other backends that rely on externally supplied mesh resources.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/1f161f706e17b2b5. Report an issue: GitHub.