{"record":{"id":"3ec6c70d27b36af6","repo":"jax-ml/jax","slug":"unsupported-a-register-array-dtype-a-mlir-dtype","errorCode":null,"errorMessage":"Unsupported A register array dtype: {a.mlir_dtype}","messagePattern":"Unsupported A register array dtype: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":157,"sourceCode":"    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):\n    if a.mlir_dtype not in {bf16, f16, i8, f8e5m2, f8e4m3fn}:\n      raise ValueError(f\"Unsupported A register array dtype: {a.mlir_dtype}\")\n    # Column count must be equal to swizzle // bytewidth.\n    elt_bytewidth = utils.bytewidth(a_element_type)\n    swizzle_elems = swizzle // elt_bytewidth\n    if a.shape != (64, swizzle_elems):\n      raise ValueError(\"Unsupported A register array shape\")\n    if a.layout not in {fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_8BIT}:\n      raise ValueError(\"Unsupported A register array layout\")\n    if a_k_stride is not None or a_transpose is not None:\n      raise ValueError(\"Unsupported WGMMA features with A in registers\")\n  else:\n    if a_k_stride is None or a_k_stride % 16:\n      raise ValueError\n    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)","sourceCodeStart":139,"sourceCodeEnd":175,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L139-L175","documentation":"At wgmma.py:157, when A is supplied as a FragmentedArray in registers, its dtype must be one of bf16/f16/i8/f8e5m2/f8e4m3fn — the types the wgmma A-register form accepts.","triggerScenarios":"Calling wgmma.wgmma with a as a FragmentedArray of any other dtype (f32, s32, u8, etc.).","commonSituations":"Keeping A in registers from a prior computation without converting to a hardware dtype; f32 activations fed directly into wgmma.","solutions":["Convert A to bf16/f16/i8/f8 before passing (a.to_dtype(...))","Or pass A as an SMEM memref instead of registers, letting the normal path validate","Ensure quantization/casting happens before the wgmma call in the pipeline"],"exampleFix":"# before\nacc = wgmma.wgmma(a_f32_regs, b, acc, ...)\n# after\na = a_f32_regs.to_dtype(ir.BF16Type.get())\nacc = wgmma.wgmma(a, b, acc, ...)","handlingStrategy":"type-guard","validationCode":"assert a.mlir_dtype in {ir.BF16Type.get(), ir.F16Type.get(), ir.IntegerType.get_signless(8)}, 'cast A before wgmma'","typeGuard":"def wgmma_a_ok(a):\n    import jax.experimental.mosaic.gpu as mgpu\n    return str(a.mlir_dtype) in ('bf16', 'f16', 'i8', 'f8E5M2', 'f8E4M3FNUZ') or a.mlir_dtype in {\n        ir.BF16Type.get(), ir.F16Type.get(), ir.IntegerType.get_signless(8),\n        ir.Float8E5M2Type.get(), ir.Float8E4M3FNType.get()}","tryCatchPattern":null,"preventionTips":["Cast activations to bf16/f16 before the wgmma stage","Keep a single dtype-conversion point in the kernel"],"tags":["jax","mosaic-gpu","wgmma","dtype","registers"],"backgroundTag":"unsupported-type-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}