jax-ml/jax · error · NotImplementedError

Mesh {mesh} is not supported by the Mosaic GPU backend

Error message

Mesh {mesh} is not supported by the Mosaic GPU backend

What it means

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.

Source

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

    input_output_aliases: tuple[tuple[int, int], ...],
    grid_mapping: pallas_core.GridMapping,
    mesh: pallas_core.Mesh | None,
    compiler_params: pallas_core.CompilerParams | None,
    cost_estimate: pallas_core.CostEstimate | None,
    out_avals: tuple[jax_core.AbstractValue, ...],
    metadata: frozen_dict.FrozenDict[str, str] | None,
    name: str | None,
):
  del metadata, name  # TODO(sharadmv): Add metadata to HLO.
  debug_info = jaxpr.debug_info
  del interpret, out_avals
  if grid_mapping.num_dynamic_grid_bounds:
    raise NotImplementedError(
        "dynamic grid bounds not supported in the Mosaic GPU backend"
    )

  if mesh is not None and not isinstance(mesh, gpu_core.Mesh):
    raise NotImplementedError(
        f"Mesh {mesh} is not supported by the Mosaic GPU backend"
    )

  if debug:
    print(f"\nThe kernel jaxpr for pallas_call {debug_info.func_src_info}:")
    print(jaxpr)
    print(f"The grid mapping for pallas_call {debug_info.func_src_info}:")
    print(grid_mapping)

  mgpu.dialect.register_dialect(ctx.module_context.context)

  if compiler_params is None:
    gpu_params = gpu_core.CompilerParams()
  else:
    assert isinstance(compiler_params, gpu_core.CompilerParams)
    gpu_params = compiler_params

  jax_mesh = None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass None for mesh, or construct the correct gpu_core.Mesh
2. Branch per backend in shared launch code

Example fix

# before
kernel[grid, mesh=jmesh](...)  # jax.sharding.Mesh
# after
kernel[grid](...)  # or pass gpu_core.Mesh instance

from jax._src.pallas import gpu_core
assert mesh is None or isinstance(mesh, gpu_core.Mesh)

def is_gpu_mesh(mesh) -> bool:
    from jax._src.pallas import gpu_core
    return mesh is None or isinstance(mesh, gpu_core.Mesh)

When it happens

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

Common situations: Sharing kernel-launch helper code between TPU and GPU backends; accidentally passing jax.sharding.Mesh where plgpu.Mesh is expected.

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
• {axis} mixes JAX mesh and Pallas mesh grid axes
• 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/25ba01100f8af477. Report an issue: GitHub.