{"record":{"id":"d450a657f247f49c","repo":"jax-ml/jax","slug":"wgmma-instruction-only-supports-f32-f16-and-s32-o","errorCode":null,"errorMessage":"WGMMA instruction only supports f32, f16 and s32 out (got {out_ty})","messagePattern":"WGMMA instruction only supports f32, f16 and s32 out \\(got (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":188,"sourceCode":"    if a_transpose is None:\n      raise ValueError\n\n  if isinstance(out_ty, ir.F32Type) or out_ty == i32:\n    num_acc_regs = n // 2\n    out_ty_field = ir.VectorType.get((1,), out_ty)\n    acc_regs = list(acc.flat)\n    assert acc_regs[0].type == ir.VectorType.get((1,), out_ty)\n    to_acc_vec_regs = lambda regs: np.array(regs).reshape(acc.shape)\n    acc_constraint = \"r\" if isinstance(out_ty, ir.IntegerType) else \"f\"\n  elif isinstance(out_ty, ir.F16Type):\n    num_acc_regs = n // 4\n    out_ty_field = i32\n    acc_regs = [_as_i32_reg(reg) for reg in acc.flat]\n    vec_ty = ir.VectorType(acc.flat[0].type)\n    to_acc_vec_regs = lambda regs: np.array([_unpack_i32(vec_ty, reg) for reg in regs]).reshape(acc.shape)\n    acc_constraint = \"r\"\n  else:\n    raise ValueError(\n        f\"WGMMA instruction only supports f32, f16 and s32 out (got {out_ty})\")\n\n  if supports_transpose:\n    num_imm_regs = 4\n  elif out_ty == i32:\n    num_imm_regs = 0\n  else:\n    num_imm_regs = 2\n\n  if a_in_regs:\n    a_reg_constraints = [\"r\"] * 4  # 4x (b)f16x2 or s8x4 registers\n    if supports_transpose:\n      num_imm_regs -= 1  # transpose not supported for a in registers\n  else:\n    a_reg_constraints = [\"l\"]  # descriptor\n  # Reference for i/o aliasing: https://gcc.gnu.org/onlinedocs/gcc/Extended-Asm.html\n  # Seems like it's not actually documented in LLVM IR docs.\n  reg_constraints_list = (","sourceCodeStart":170,"sourceCodeEnd":206,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L170-L206","documentation":"At wgmma.py:188, after the accumulator dtype branch, any out_ty other than f32, f16, or s32 (i32) is rejected — the wgmma instruction's accumulator register file only holds those types.","triggerScenarios":"Calling wgmma.wgmma with an accumulator FragmentedArray of dtype f64, bf16, s64, etc. (this check complements _supported_wgmma_types).","commonSituations":"Creating the accumulator with a dtype matched to the operands instead of the output; bf16 accumulation attempts; custom vector types on acc.flat.","solutions":["Allocate the accumulator as f32 (default), f16, or signless i32 (WGMMAAccumulator.zero does this correctly)","Convert results to the desired dtype after the wgmma loop","Check that acc was built from WGMMAAccumulator.zero or a prior wgmma result"],"exampleFix":"# before\nacc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.BF16Type.get())\n# after\nacc = wgmma.WGMMAAccumulator.zero(64, 64, dtype=ir.F32Type.get())\nresult_bf16 = acc.value.to_dtype(ir.BF16Type.get())","handlingStrategy":"validation","validationCode":"assert str(out_ty) in ('f32', 'f16', 'i32'), f'bad accumulator dtype {out_ty}'","typeGuard":"def wgmma_out_ok(dt): return str(dt) in ('f32', 'f16', 'i32')","tryCatchPattern":null,"preventionTips":["Create accumulators only via WGMMAAccumulator.zero","Cast after accumulation, not before"],"tags":["jax","mosaic-gpu","wgmma","dtype","accumulator"],"backgroundTag":"unsupported-type-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}