{"record":{"id":"26adc9e56fab3b03","repo":"jax-ml/jax","slug":"inline-mgpu-p-does-not-support-discharge","errorCode":null,"errorMessage":"inline_mgpu_p does not support discharge.","messagePattern":"inline_mgpu_p does not support discharge\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3519,"sourceCode":"      jax_core.ShapedArray(x.shape, x.dtype) for x in flat_ret_ty\n  )\n  # TODO(cperivol): Let the user set the effects.\n  flat_args = flat_args_and_transforms[:pytree_args.num_leaves]\n  return aval_return, {\n      gpu_core._wgmma_pipeline_effect,\n      gpu_core._memory_effect,\n      *itertools.chain.from_iterable(\n          (state.ReadEffect(i), state.WriteEffect(i))\n          for i, r in enumerate(flat_args)\n          if isinstance(r, state.AbstractRef)\n      ),\n  }\n\n\n@discharge.register_discharge_rule(inline_mgpu_p)\ndef _inline_mgpu_discharge(*args, **kwargs):\n  del args, kwargs\n  raise NotImplementedError(\"inline_mgpu_p does not support discharge.\")\n\n\ndef _type_check_mgpu_lane_semantics(v, ty):\n  match (ty, v):\n    case (RefType(), ir.Value()) if isinstance(v.type, ir.MemRefType):\n      pass\n    case (ShapeDtypeStruct(), mgpu.FragmentedArray()):\n      mlir_dtype = mgpu_utils.dtype_to_ir_type(ty.dtype)\n      if v.mlir_dtype != mlir_dtype:\n        raise ValueError(\n            f\"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}.\"\n        )\n      if ty.shape != v.shape:\n        raise ValueError(\n            f\"Array shape mismatch: expected {ty.shape} got {v.shape}.\"\n        )\n      if v.layout != ty.layout.to_mgpu():\n        raise ValueError(","sourceCodeStart":3501,"sourceCodeEnd":3537,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3501-L3537","documentation":"inline_mgpu primitives perform raw MLIR operations inside a Pallas kernel and have no pure-JAX fallback; attempting to discharge (lower a Pallas kernel to plain JAX, e.g. for CPU interpretation or pipelining outside GPU) is explicitly unsupported.","triggerScenarios":"Running a kernel containing inline_mgpu/inline_ptx through discharge-based paths: interpretation on CPU, pallas discharge mode, or APIs that require discharging the kernel.","commonSituations":"Trying to debug a Pallas kernel with the interpreter; using a pipeline that discharges transforms (e.g. certain checkpointing/remat paths).","solutions":["Run the kernel on GPU (compiled path) instead of the interpreter/discharge path","Replace inline_mgpu usage with standard Pallas/Mosaic primitives that support discharge when interpretation is required"],"exampleFix":"# before\nkernel_with_inline_mgpu(...)  # interpreted/discharge run\n# after\ncompiled_gpu_kernel(...)  # run via pallas_call with a GPU compilation target","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    discharge_based_run(kernel)\nexcept NotImplementedError:\n    run_on_gpu(kernel)  # compiled path","preventionTips":["Avoid inline_mgpu in kernels that need interpretation","Provide a non-inline fallback kernel for CPU debugging"],"tags":["jax","pallas","inline-mgpu","discharge","not-implemented"],"backgroundTag":"operation-not-supported-in-fallback-mode","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}