jax-ml/jax · error · NotImplementedError

Scratch operands are not supported by the Mosaic GPU mpmd_ma

Error message

Scratch operands are not supported by the Mosaic GPU mpmd_map lowering

What it means

mpmd_map on GPU validates that jaxpr.invars count equals inputs + outputs; extra invars imply scratch operands, which the GPU backend does not support because scratch is handled internally by plgpu.kernel.

Source

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

    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}:")
    print(jaxpr)

  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
  axis_context = ctx.module_context.axis_context
  if axis_context is not None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove scratch buffers from the mpmd_map function signature; allocate scratch inside the kernel via plgpu mechanisms
2. Refactor to plain pallas_call on GPU

assert len(jaxpr.invars) == len(inputs) + len(outputs), "scratch operands unsupported on GPU mpmd_map"

When it happens

Trigger: Supplying a jaxpr to mpmd_map that has scratch (intermediate workspace) operands in its signature when lowering for Mosaic GPU.

Common situations: Kernels written for TPU mpmd_map that declare scratch buffers explicitly.

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

• Lowering multiple mesh/function pairs is not supported by th
• External meshes are not supported by the Mosaic GPU backend
• 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/4f28d7a96a164bcc. Report an issue: GitHub.