jax-ml/jax · error · NotImplementedError

Lowering multiple mesh/function pairs is not supported by th

Error message

Lowering multiple mesh/function pairs is not supported by the Mosaic GPU backend

What it means

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.

Source

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

def mpmd_map_mgpu_lowering_rule(
    ctx: mlir.LoweringRuleContext,
    *args,
    meshes,
    jaxprs,
    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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split the mpmd_map into separate single-function kernel launches
2. Await or contribute upstream support for multi-pair mpmd_map on GPU

assert len(jaxprs) == 1, "mpmd_map on GPU supports exactly one function pair"

When it happens

Trigger: 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.

Common situations: Reusing MPMD pipelines written for TPU on GPU; building multi-kernel launch configurations.

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

• External meshes are not supported by the Mosaic GPU backend
• Scratch operands are not supported by the Mosaic GPU mpmd_ma
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}

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