{"record":{"id":"f6dc091d25883621","repo":"jax-ml/jax","slug":"accumulator-and-rhs-have-incompatible-shapes-expe-f6dc09","errorCode":null,"errorMessage":"Accumulator and RHS have incompatible shapes. Expected RHS to have shape (k, n) and accumulator to have shape (m, n). Accumulator: {acc.shape}. RHS: {b.shape}.","messagePattern":"Accumulator and RHS have incompatible shapes\\. Expected RHS to have shape \\(k, n\\) and accumulator to have shape \\(m, n\\)\\. Accumulator: (.+?)\\. RHS: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2458,"sourceCode":"  rhs_k, rhs_n = b.shape\n  is_sparse = a_sparse_metadata is not None\n\n  if acc_m != lhs_m:\n    raise ValueError(\n        \"Accumulator and LHS have incompatible shapes. Expected LHS to have\"\n        \" shape (m, k) and accumulator to have shape (m, n). Accumulator:\"\n        f\" {acc.shape}. LHS: {a.shape}.\"\n    )\n\n  if collective_axis is not None:\n    if acc_n != rhs_n * 2:\n      raise ValueError(\n          \"Accumulator and RHS have incompatible shapes. Expected RHS to have \"\n          \"shape (k, n) and accumulator to have shape (m, n * 2) in \"\n          f\"collective mode. Accumulator: {acc.shape}. RHS: {b.shape}.\"\n      )\n  elif acc_n != rhs_n:\n    raise ValueError(\n        \"Accumulator and RHS have incompatible shapes. Expected RHS to have\"\n        \" shape (k, n) and accumulator to have shape (m, n). Accumulator:\"\n        f\" {acc.shape}. RHS: {b.shape}.\"\n    )\n\n  if (lhs_k * (1 + is_sparse)) != rhs_k:\n    raise ValueError(\n        f\"LHS and RHS have incompatible shapes. LHS: {a.shape}. RHS: {b.shape}.\")\n\n  if isinstance(acc, pallas_core.TransformedRef):\n    acc_transforms_leaves, acc_transforms_tree = jax.tree.flatten(\n        acc.transforms)\n    acc = acc.ref\n  else:\n    acc_transforms_leaves, acc_transforms_tree = [], None\n\n  if isinstance(a, pallas_core.TransformedRef):\n    a_transforms_leaves, a_transforms_tree = jax.tree.flatten(a.transforms)","sourceCodeStart":2440,"sourceCodeEnd":2476,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2440-L2476","documentation":"In non-collective tcgen05.mma (collective_axis=None), the accumulator's N must equal the RHS's n: acc.shape[1] != b.shape[1].","triggerScenarios":"Calling tcgen05.mma with accumulator (m, 128) but RHS block of shape (k, 64), or vice versa.","commonSituations":"Changing the N tiling of the B matrix without updating the accumulator allocation; copying accumulator sizing from a collective-mode kernel into a non-collective one (off-by-2x N).","solutions":["Match accumulator N to b.shape[1] exactly","If porting from collective mode, halve the accumulator N","Validate all three shapes (acc (m,n), a (m,k), b (k,n)) before the call"],"exampleFix":"# before\nacc = allocate(TMEM, (m, 2 * n), jnp.float32)  # copied from collective kernel\ntcgen05.mma(a, b, acc, k_dim=k)\n# after\nacc = allocate(TMEM, (m, n), jnp.float32)\ntcgen05.mma(a, b, acc, k_dim=k)","handlingStrategy":"validation","validationCode":"assert acc.shape[1] == b.shape[1], (acc.shape, b.shape)\ntcgen05.mma(a, b, acc, k_dim=a.shape[1])","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Add a shape-check helper used by all tcgen05.mma call sites","Avoid copying sizes between collective and non-collective kernels"],"tags":["jax","pallas","tcgen05","mma","shape-mismatch"],"backgroundTag":"matmul-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}