jax-ml/jax · error · ValueError

Unsupported A register array shape

Error message

Unsupported A register array shape

What it means

At wgmma.py:162, an A FragmentedArray passed in registers must have shape exactly (64, swizzle // bytewidth(element_type)) — the register footprint wgmma expects for one instruction.

Source

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

  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)
    acc_regs = list(acc.flat)
    assert acc_regs[0].type == ir.VectorType.get((1,), out_ty)
    to_acc_vec_regs = lambda regs: np.array(regs).reshape(acc.shape)
    acc_constraint = "r" if isinstance(out_ty, ir.IntegerType) else "f"
  elif isinstance(out_ty, ir.F16Type):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Resize A so its columns equal swizzle // bytewidth(a_element_type)
2. Adjust swizzle to match the A tile width (must be one of 32/64/128)
3. Loop over K in chunks, calling wgmma per (64, swizzle//bytewidth) slice of A

Example fix

# before
a = a.reshape(64, 128)
acc = wgmma.wgmma(a, b, acc, swizzle=64)  # f16: expects (64, 32)
# after
acc = wgmma.wgmma(a[:, :32], b, acc, swizzle=64)

expected_cols = swizzle // utils.bytewidth(a_element_type)
assert a.shape == (64, expected_cols)

When it happens

Trigger: Calling wgmma.wgmma with an in-register A whose shape is e.g. (64, 128) while swizzle=64 and dtype is f16 (expected (64, 32)).

Common situations: Mismatching the swizzle parameter and the A tile width; using a k-dim tile larger than one wgmma instruction without splitting; dtype changes (i8 doubles the element count).

Related errors

• No swizzle is not supported
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• Swizzle {self.swizzle} requires the trailing dimension to be
• Can't transpose the swizzled dimension.
• Reshape shape {shape} is not divisible by swizzle elements {

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