{"record":{"id":"142a88e93c6d6331","repo":"jax-ml/jax","slug":"accumulator-and-rhs-have-incompatible-shapes-expe","errorCode":null,"errorMessage":"Accumulator and RHS have incompatible shapes. Expected RHS to have shape (k, n) and accumulator to have shape (m, n * 2) in collective mode. 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 \\* 2\\) in collective mode\\. Accumulator: (.+?)\\. RHS: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2452,"sourceCode":"    collective_axis: The name of the cluster axis along which to perform\n      a collective MMA. The cluster axis should have a size of exactly 2,\n      and must be on the minormost cluster axis.\n  \"\"\"\n  acc_m, acc_n = acc.shape\n  lhs_m, lhs_k = a.shape\n  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)","sourceCodeStart":2434,"sourceCodeEnd":2470,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2434-L2470","documentation":"When collective_axis is set on tcgen05.mma, the collective MMA produces an N dimension twice the RHS's n, so the accumulator must have shape (m, n*2). This fires when acc.shape[1] != b.shape[1] * 2.","triggerScenarios":"Using tcgen05.mma(..., collective_axis=0) with an accumulator sized (m, n) matching b's n instead of (m, 2*n).","commonSituations":"Switching a kernel from non-collective to collective MMA without resizing the accumulator; porting examples that omit the doubled N in collective mode.","solutions":["Allocate/size the accumulator N as 2 * b.shape[1] when using collective_axis","Double-check the collective MMA docs for the (m, n*2) accumulator layout","If you don't need 2CTA tensor-core mode, drop collective_axis"],"exampleFix":"# before\nacc = allocate(TMEM, (m, b.shape[1]), jnp.float32)\ntcgen05.mma(a, b, acc, k_dim=k, collective_axis=0)\n# after\nacc = allocate(TMEM, (m, 2 * b.shape[1]), jnp.float32)\ntcgen05.mma(a, b, acc, k_dim=k, collective_axis=0)","handlingStrategy":"validation","validationCode":"n_acc = 2 * b.shape[1] if collective_axis is not None else b.shape[1]\nassert acc.shape == (a.shape[0], n_acc)\ntcgen05.mma(a, b, acc, k_dim=a.shape[1], collective_axis=collective_axis)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Wrap tcgen05.mma in a helper that computes the expected accumulator shape","Document the 2x N rule for collective mode"],"tags":["jax","pallas","tcgen05","collective-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"}