{"record":{"id":"07d1b4ff530e0f4d","repo":"jax-ml/jax","slug":"unsupported-wgmma-types-out-ty-a-element-type","errorCode":null,"errorMessage":"Unsupported wgmma types {(out_ty, a_element_type)=}","messagePattern":"Unsupported wgmma types (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":135,"sourceCode":"    return False\n\n\ndef wgmma_m64(\n    acc: np.ndarray,  # of register Values\n    a,\n    b_descriptor: ir.Value,\n    a_transpose: bool | None,\n    b_transpose: bool,\n    a_k_stride: int | None,\n    b_k_stride: int,\n    n: int,\n    swizzle: int,\n    a_element_type: ir.Type,\n    b_element_type: ir.Type,\n):\n  out_ty = ir.VectorType(acc.flat[0].type).element_type\n  if not _supported_wgmma_types(out_ty, a_element_type):\n    raise ValueError(f\"Unsupported wgmma types {(out_ty, a_element_type)=}\")\n  if not _supported_wgmma_types(out_ty, b_element_type):\n    raise ValueError(f\"Unsupported wgmma types {(out_ty, b_element_type)=}\")\n  if n % 8:\n    raise ValueError\n\n  bf16 = ir.BF16Type.get()\n  f16 = ir.F16Type.get()\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  i64 = ir.IntegerType.get_signless(64)\n  f8e5m2 = ir.Float8E5M2Type.get()\n  f8e4m3fn = ir.Float8E4M3FNType.get()\n  if b_k_stride % 16:\n    raise ValueError\n  assert bytewidth(a_element_type) == bytewidth(b_element_type)\n  # Only 16-bit types support transposes\n  supports_transpose = bytewidth(b_element_type) == 2\n  if not supports_transpose and (a_transpose or b_transpose):","sourceCodeStart":117,"sourceCodeEnd":153,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L117-L153","documentation":"wgmma_m64 (wgmma.py:135) validates that (accumulator element type, A operand element type) is one of the combinations the hardware supports (f32/f16/s32 accumulators over f16/bf16/i8/f8 operand pairs). Unsupported pairs raise this error.","triggerScenarios":"Calling wgmma.wgmma(...) with e.g. an f64 or bf16 accumulator, or an A operand dtype (e.g. f32) that cannot feed the tensor core for the chosen accumulator type.","commonSituations":"Writing generic GEMM kernels that reuse one dtype for everything; f32 A operands; mixing an s8 A with an f16 accumulator.","solutions":["Use f16/bf16/i8/f8e5m2/f8e4m3fn for A and f32 (or matching f16/s32) for the accumulator","If operands are f32, convert them to bf16/f16 before the wgmma call","Check _supported_wgmma_types in the same file for your JAX version's exact allowed pairs"],"exampleFix":"# before\nacc = wgmma.wgmma(a_f32_memref, b_f32_memref, acc_f32)\n# after\na = a.to_dtype(ir.BF16Type.get())\nb = b.to_dtype(ir.BF16Type.get())\nacc = wgmma.wgmma(a, b, acc_f32)","handlingStrategy":"type-guard","validationCode":"assert (str(out_ty), str(a_ty)) in SUPPORTED, 'unsupported wgmma type pair'  # mirror _supported_wgmma_types","typeGuard":"def supported_a(out_ty, a_ty):\n    ok = {('f32','f16'),('f32','bf16'),('f32','i8'),('f16','f16'),('i32','i8')}\n    return (str(out_ty), str(a_ty)) in ok","tryCatchPattern":null,"preventionTips":["Keep operand dtypes in {f16,bf16,i8,f8} and accumulators in {f32,f16,s32}","Centralize dtype policy for tensor-core kernels"],"tags":["jax","mosaic-gpu","wgmma","dtype","tensor-core"],"backgroundTag":"unsupported-type-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}