{"record":{"id":"26c6d584237754e8","repo":"jax-ml/jax","slug":"unsupported-a-register-array-shape","errorCode":null,"errorMessage":"Unsupported A register array shape","messagePattern":"Unsupported A register array shape","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":162,"sourceCode":"  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)\n    acc_regs = list(acc.flat)\n    assert acc_regs[0].type == ir.VectorType.get((1,), out_ty)\n    to_acc_vec_regs = lambda regs: np.array(regs).reshape(acc.shape)\n    acc_constraint = \"r\" if isinstance(out_ty, ir.IntegerType) else \"f\"\n  elif isinstance(out_ty, ir.F16Type):","sourceCodeStart":144,"sourceCodeEnd":180,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L144-L180","documentation":"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.","triggerScenarios":"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)).","commonSituations":"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).","solutions":["Resize A so its columns equal swizzle // bytewidth(a_element_type)","Adjust swizzle to match the A tile width (must be one of 32/64/128)","Loop over K in chunks, calling wgmma per (64, swizzle//bytewidth) slice of A"],"exampleFix":"# before\na = a.reshape(64, 128)\nacc = wgmma.wgmma(a, b, acc, swizzle=64)  # f16: expects (64, 32)\n# after\nacc = wgmma.wgmma(a[:, :32], b, acc, swizzle=64)","handlingStrategy":"validation","validationCode":"expected_cols = swizzle // utils.bytewidth(a_element_type)\nassert a.shape == (64, expected_cols)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive A tile width from swizzle // bytewidth, not independently","Split large K tiles into per-instruction chunks"],"tags":["jax","mosaic-gpu","wgmma","shape","swizzle"],"backgroundTag":"gpu-tile-shape-invalid","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}