{"record":{"id":"ea4fc60343044759","repo":"jax-ml/jax","slug":"k-mismatch-k-k2","errorCode":null,"errorMessage":"K mismatch: {k} != {k2}","messagePattern":"K mismatch: (.+?) != (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/mma.py","lineNumber":196,"sourceCode":"    a: A `FragmentedArray` with a `TiledLayout`  generated from\n      `MMALayouts.lhs`.\n    b: A `FragmentedArray` with a `TiledLayout` generated from `MMALayouts.rhs`.\n\n  Returns:\n    A new `FragmentedArray` with the result of the computation with\n      the same type as `acc`.\n  \"\"\"\n\n  (m, k) = a.shape\n  (k2, n) = b.shape\n  (m2, n2) = acc.shape\n\n  if m != m2:\n    raise ValueError(f\"M mismatch: {m} != {m2}\")\n  if n != n2:\n    raise ValueError(f\"N mismatch: {n} != {n2}\")\n  if k != k2:\n    raise ValueError(f\"K mismatch: {k} != {k2}\")\n\n  # todo(cperivol): A tile shape can have dimensions that are higher\n  # multiples of the mma op size as long as those dimensions are not\n  # sharded across warps.\n  i4 = ir.IntegerType.get_signless(4)\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  bf16 = ir.BF16Type.get()\n  f16 = ir.F16Type.get()\n  f8e4m3fn = ir.Float8E4M3FNType.get()\n  f8e5m2 = ir.Float8E5M2Type.get()\n  if (element_type := a.mlir_dtype) != b.mlir_dtype:\n    raise ValueError(f\"Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}\")\n  if element_type not in (bf16, f16, f8e4m3fn, f8e5m2, i8, i4):\n    raise NotImplementedError(f\"Unsupported operand type: {element_type}\")\n  if isinstance(element_type, ir.IntegerType):\n    if acc.mlir_dtype != i32:\n      raise NotImplementedError(\"Only s32 accumulator supported for integer operands.\")","sourceCodeStart":178,"sourceCodeEnd":214,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/mma.py#L178-L214","documentation":"The contraction dimension must agree: a.shape[1] must equal b.shape[0]. A K mismatch means the operand tiles cannot be multiplied.","triggerScenarios":"Calling mma(a, b, acc) where a's K dim differs from b's K dim, e.g. slicing one operand differently along the contraction axis.","commonSituations":"Split-K or partial-accumulation loops slicing a and b inconsistently; off-by-one loop bounds over K.","solutions":["Slice a and b with the same K extent in each mma call","Check loop bounds for the contraction dimension","Re-verify tile shapes after reshaping operands"],"exampleFix":"// before\nacc = mma.mma(a[:, :32], b[64:, :], acc)\n// after\nacc = mma.mma(a[:, :32], b[:32, :], acc)","handlingStrategy":"validation","validationCode":"assert a.shape[1] == b.shape[0], 'K mismatch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Slice contraction dims symmetrically in split-K loops"],"tags":["jax","mosaic","mma","shape-mismatch"],"backgroundTag":"matrix-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}