jax-ml/jax · error · ValueError

Unsupported A register array dtype: {a.mlir_dtype}

Error message

Unsupported A register array dtype: {a.mlir_dtype}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/wgmma.py:157

    raise ValueError

  bf16 = ir.BF16Type.get()
  f16 = ir.F16Type.get()
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)
  i64 = ir.IntegerType.get_signless(64)
  f8e5m2 = ir.Float8E5M2Type.get()
  f8e4m3fn = ir.Float8E4M3FNType.get()
  if b_k_stride % 16:
    raise ValueError
  assert bytewidth(a_element_type) == bytewidth(b_element_type)
  # Only 16-bit types support transposes
  supports_transpose = bytewidth(b_element_type) == 2
  if not supports_transpose and (a_transpose or b_transpose):
    raise ValueError("Only f16 WGMMA supports transposes")
  if a_in_regs := isinstance(a, fa.FragmentedArray):
    if a.mlir_dtype not in {bf16, f16, i8, f8e5m2, f8e4m3fn}:
      raise ValueError(f"Unsupported A register array dtype: {a.mlir_dtype}")
    # Column count must be equal to swizzle // bytewidth.
    elt_bytewidth = utils.bytewidth(a_element_type)
    swizzle_elems = swizzle // elt_bytewidth
    if a.shape != (64, swizzle_elems):
      raise ValueError("Unsupported A register array shape")
    if a.layout not in {fa.WGMMA_LAYOUT, fa.WGMMA_LAYOUT_8BIT}:
      raise ValueError("Unsupported A register array layout")
    if a_k_stride is not None or a_transpose is not None:
      raise ValueError("Unsupported WGMMA features with A in registers")
  else:
    if a_k_stride is None or a_k_stride % 16:
      raise ValueError
    if a_transpose is None:
      raise ValueError

  if isinstance(out_ty, ir.F32Type) or out_ty == i32:
    num_acc_regs = n // 2
    out_ty_field = ir.VectorType.get((1,), out_ty)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert A to bf16/f16/i8/f8 before passing (a.to_dtype(...))
2. Or pass A as an SMEM memref instead of registers, letting the normal path validate
3. Ensure quantization/casting happens before the wgmma call in the pipeline

Example fix

# before
acc = wgmma.wgmma(a_f32_regs, b, acc, ...)
# after
a = a_f32_regs.to_dtype(ir.BF16Type.get())
acc = wgmma.wgmma(a, b, acc, ...)

assert a.mlir_dtype in {ir.BF16Type.get(), ir.F16Type.get(), ir.IntegerType.get_signless(8)}, 'cast A before wgmma'

def wgmma_a_ok(a):
    import jax.experimental.mosaic.gpu as mgpu
    return str(a.mlir_dtype) in ('bf16', 'f16', 'i8', 'f8E5M2', 'f8E4M3FNUZ') or a.mlir_dtype in {
        ir.BF16Type.get(), ir.F16Type.get(), ir.IntegerType.get_signless(8),
        ir.Float8E5M2Type.get(), ir.Float8E4M3FNType.get()}

When it happens

Trigger: Calling wgmma.wgmma with a as a FragmentedArray of any other dtype (f32, s32, u8, etc.).

Common situations: Keeping A in registers from a prior computation without converting to a hardware dtype; f32 activations fed directly into wgmma.

Related errors

• Unsupported wgmma types {(out_ty, a_element_type)=}
• Unsupported wgmma types {(out_ty, b_element_type)=}
• Only f16 WGMMA supports transposes
• Unsupported WGMMA features with A in registers
• WGMMA instruction only supports f32, f16 and s32 out (got {o

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/3ec6c70d27b36af6. Report an issue: GitHub.