jax-ml/jax · error · ValueError

Unsupported wgmma types {(out_ty, a_element_type)=}

Error message

Unsupported wgmma types {(out_ty, a_element_type)=}

What it means

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.

Source

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

    return False


def wgmma_m64(
    acc: np.ndarray,  # of register Values
    a,
    b_descriptor: ir.Value,
    a_transpose: bool | None,
    b_transpose: bool,
    a_k_stride: int | None,
    b_k_stride: int,
    n: int,
    swizzle: int,
    a_element_type: ir.Type,
    b_element_type: ir.Type,
):
  out_ty = ir.VectorType(acc.flat[0].type).element_type
  if not _supported_wgmma_types(out_ty, a_element_type):
    raise ValueError(f"Unsupported wgmma types {(out_ty, a_element_type)=}")
  if not _supported_wgmma_types(out_ty, b_element_type):
    raise ValueError(f"Unsupported wgmma types {(out_ty, b_element_type)=}")
  if n % 8:
    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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use f16/bf16/i8/f8e5m2/f8e4m3fn for A and f32 (or matching f16/s32) for the accumulator
2. If operands are f32, convert them to bf16/f16 before the wgmma call
3. Check _supported_wgmma_types in the same file for your JAX version's exact allowed pairs

Example fix

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

assert (str(out_ty), str(a_ty)) in SUPPORTED, 'unsupported wgmma type pair'  # mirror _supported_wgmma_types

def supported_a(out_ty, a_ty):
    ok = {('f32','f16'),('f32','bf16'),('f32','i8'),('f16','f16'),('i32','i8')}
    return (str(out_ty), str(a_ty)) in ok

When it happens

Trigger: 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.

Common situations: Writing generic GEMM kernels that reuse one dtype for everything; f32 A operands; mixing an s8 A with an f16 accumulator.

Related errors

• Unsupported wgmma types {(out_ty, b_element_type)=}
• Only f16 WGMMA supports transposes
• Unsupported A register array dtype: {a.mlir_dtype}
• WGMMA instruction only supports f32, f16 and s32 out (got {o
• Unsupported dtype: {ref.dtype}

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