{"record":{"id":"4f28d7a96a164bcc","repo":"jax-ml/jax","slug":"scratch-operands-are-not-supported-by-the-mosaic-g","errorCode":null,"errorMessage":"Scratch operands are not supported by the Mosaic GPU mpmd_map lowering","messagePattern":"Scratch operands are not supported by the Mosaic GPU mpmd_map lowering","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/pallas_call_registration.py","lineNumber":217,"sourceCode":"    raise NotImplementedError(\n        \"Lowering multiple mesh/function pairs is not supported by the Mosaic\"\n        \" GPU backend\"\n    )\n  if external_meshes:\n    raise NotImplementedError(\n        \"External meshes are not supported by the Mosaic GPU backend\"\n    )\n  [jaxpr] = jaxprs\n  [mesh] = meshes\n  if not isinstance(mesh, gpu_core.Mesh):\n    raise NotImplementedError(\n        f\"Mesh {mesh} is not supported by the Mosaic GPU backend\"\n    )\n  # On GPU ``mpmd_map`` kernels never carry scratch operands -- scratch is\n  # handled separately by ``plgpu.kernel``. So the jaxpr invars are exactly the\n  # inputs followed by the outputs.\n  if len(jaxpr.invars) != len(args) + len(ctx.avals_out):\n    raise NotImplementedError(\n        \"Scratch operands are not supported by the Mosaic GPU mpmd_map lowering\"\n    )\n\n  if debug:\n    print(f\"\\nThe kernel jaxpr for mpmd_map {jaxpr.debug_info.func_src_info}:\")\n    print(jaxpr)\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\n  axis_context = ctx.module_context.axis_context\n  if axis_context is not None:","sourceCodeStart":199,"sourceCodeEnd":235,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py#L199-L235","documentation":"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.","triggerScenarios":"Supplying a jaxpr to mpmd_map that has scratch (intermediate workspace) operands in its signature when lowering for Mosaic GPU.","commonSituations":"Kernels written for TPU mpmd_map that declare scratch buffers explicitly.","solutions":["Remove scratch buffers from the mpmd_map function signature; allocate scratch inside the kernel via plgpu mechanisms","Refactor to plain pallas_call on GPU"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert len(jaxpr.invars) == len(inputs) + len(outputs), \"scratch operands unsupported on GPU mpmd_map\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't declare scratch in mpmd_map signatures on GPU"],"tags":["pallas","mosaic-gpu","mpmd-map","scratch","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}