{"record":{"id":"def2afe16683c30d","repo":"jax-ml/jax","slug":"n-mismatch-n-n2","errorCode":null,"errorMessage":"N mismatch: {n} != {n2}","messagePattern":"N mismatch: (.+?) != (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/mma.py","lineNumber":194,"sourceCode":"    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):\n    raise NotImplementedError(f\"Unsupported operand type: {element_type}\")\n  if isinstance(element_type, ir.IntegerType):","sourceCodeStart":176,"sourceCodeEnd":212,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/mma.py#L176-L212","documentation":"The N dimension of the accumulator must equal b.shape[1] (acc.shape[1] == b.shape[1]). A mismatch indicates the matmul tile's output width doesn't line up with operand b.","triggerScenarios":"Calling mma(a, b, acc) with an accumulator whose second dim differs from b's second dim, e.g. N-tile changed without resizing acc.","commonSituations":"Sweeping N tile sizes while reusing a stale accumulator buffer, or transposing b without adjusting acc.","solutions":["Allocate/reallocate acc with shape (m, b.shape[1])","Keep N of b and acc consistent when retuning tile sizes"],"exampleFix":"// before\nacc = fa.from_tensor(jnp.zeros((m, 64), jnp.float32))\nacc = mma.mma(a, b_n128, acc)\n// after\nacc = fa.from_tensor(jnp.zeros((m, 128), jnp.float32))\nacc = mma.mma(a, b_n128, acc)","handlingStrategy":"validation","validationCode":"assert b.shape[1] == acc.shape[1], 'N mismatch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Reallocate acc when tile N changes; derive acc shape from operands, not constants"],"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"}