{"record":{"id":"b161cd4e41ae5a72","repo":"jax-ml/jax","slug":"m-mismatch-m-m2","errorCode":null,"errorMessage":"M mismatch: {m} != {m2}","messagePattern":"M mismatch: (.+?) != (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/mma.py","lineNumber":192,"sourceCode":"\n  Args:\n    acc: A `FragmentedArray` with a `TiledLayout` generated from\n      `MMALayouts.acc`.\n    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):","sourceCodeStart":174,"sourceCodeEnd":210,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/mma.py#L174-L210","documentation":"The low-level mma helper requires the M dimension of the accumulator to match the M of operand a (a.shape[0] == acc.shape[0]). A mismatch means the shapes don't form a valid matmul tile.","triggerScenarios":"Calling mma(a, b, acc) where acc has fewer/more rows than a, e.g. acc of shape (32, N) with a of shape (64, K).","commonSituations":"Accumulator allocated with a different tile shape than the operands, or reusing an accumulator across matmuls of different M.","solutions":["Allocate acc with shape (a.shape[0], b.shape[1])","Fix the operand tiling so M matches the accumulator","Zero-init a correctly shaped accumulator each matmul"],"exampleFix":"// before\nacc = fa.from_tensor(jnp.zeros((32, n), jnp.float32))\nacc = mma.mma(a_64xk, b, acc)\n// after\nacc = fa.from_tensor(jnp.zeros((64, n), jnp.float32))\nacc = mma.mma(a_64xk, b, acc)","handlingStrategy":"validation","validationCode":"assert a.shape[0] == acc.shape[0], 'M mismatch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate accumulators with shape (a.shape[0], b.shape[1]) at kernel start"],"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"}