jax-ml/jax · error · ValueError

Unsupported wgmma types {(out_ty, b_element_type)=}

Error message

Unsupported wgmma types {(out_ty, b_element_type)=}

What it means

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.

Source

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

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):
    raise ValueError("Only f16 WGMMA supports transposes")
  if a_in_regs := isinstance(a, fa.FragmentedArray):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match B's dtype to the wgmma-supported set (f16/bf16/i8/f8e5m2/f8e4m3fn) relative to the accumulator
2. Cast B before the call or allocate the TMA buffer in the supported dtype
3. Inspect _supported_wgmma_types for the allowed pairs in your JAX version

Example fix

# before
acc = wgmma.wgmma(a_bf16, b_f32, acc_f32)
# after
b_bf16 = b.to_dtype(ir.BF16Type.get())
acc = wgmma.wgmma(a_bf16, b_bf16, acc_f32)

assert (str(out_ty), str(b_ty)) in SUPPORTED, 'unsupported wgmma type pair'

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

When it happens

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

Common situations: Loading B from TMA with a default f32 buffer; quantized A with non-quantized B; mismatched low-precision formats (f8 vs f16).

Related errors

• Unsupported wgmma types {(out_ty, a_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/f9f2f0377659fd1c. Report an issue: GitHub.