{"record":{"id":"f9f2f0377659fd1c","repo":"jax-ml/jax","slug":"unsupported-wgmma-types-out-ty-b-element-type","errorCode":null,"errorMessage":"Unsupported wgmma types {(out_ty, b_element_type)=}","messagePattern":"Unsupported wgmma types (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":137,"sourceCode":"\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):\n    raise ValueError(\"Only f16 WGMMA supports transposes\")\n  if a_in_regs := isinstance(a, fa.FragmentedArray):","sourceCodeStart":119,"sourceCodeEnd":155,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L119-L155","documentation":"Same validation as 4589 but for the B operand: wgmma_m64 (wgmma.py:137) checks (accumulator type, B element type) against _supported_wgmma_types and rejects combos the WGMMA unit cannot execute.","triggerScenarios":"Calling wgmma.wgmma with a B memref whose element type is unsupported for the accumulator (e.g. B in f32 with an f32 acc, or s8 B with an f16 acc).","commonSituations":"Loading B from TMA with a default f32 buffer; quantized A with non-quantized B; mismatched low-precision formats (f8 vs f16).","solutions":["Match B's dtype to the wgmma-supported set (f16/bf16/i8/f8e5m2/f8e4m3fn) relative to the accumulator","Cast B before the call or allocate the TMA buffer in the supported dtype","Inspect _supported_wgmma_types for the allowed pairs in your JAX version"],"exampleFix":"# before\nacc = wgmma.wgmma(a_bf16, b_f32, acc_f32)\n# after\nb_bf16 = b.to_dtype(ir.BF16Type.get())\nacc = wgmma.wgmma(a_bf16, b_bf16, acc_f32)","handlingStrategy":"type-guard","validationCode":"assert (str(out_ty), str(b_ty)) in SUPPORTED, 'unsupported wgmma type pair'","typeGuard":"def supported_b(out_ty, b_ty):\n    ok = {('f32','f16'),('f32','bf16'),('f32','i8'),('f16','f16'),('i32','i8')}\n    return (str(out_ty), str(b_ty)) in ok","tryCatchPattern":null,"preventionTips":["Allocate TMA buffers directly in the compute dtype","Validate both operands' dtypes in kernel prologue"],"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"}