{"record":{"id":"074017cf68dff23e","repo":"jax-ml/jax","slug":"accumulator-and-lhs-have-incompatible-shapes-expe","errorCode":null,"errorMessage":"Accumulator and LHS have incompatible shapes. Expected LHS to have shape (m, k) and accumulator to have shape (m, n). Accumulator: {acc.shape}. LHS: {a.shape}.","messagePattern":"Accumulator and LHS have incompatible shapes\\. Expected LHS to have shape \\(m, k\\) and accumulator to have shape \\(m, n\\)\\. Accumulator: (.+?)\\. LHS: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2444,"sourceCode":"      Must have orders_tensor_core set to True. If not specified, the MMA\n      completion should be explicitly observed by calling\n      :func:`jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive`\n    a_scale: An optional scale for the ``a`` operand. Must be a TMEM Ref if present.\n    b_scale: An optional scale for the ``b`` operand. Must be a TMEM Ref if present.\n    a_sparse_metadata: An optional sparse metadata for the ``a`` operand.\n      Must be a TMEM Ref if present.\n    accumulate: Whether to accumulate into acc or overwrite it.\n    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    )","sourceCodeStart":2426,"sourceCodeEnd":2462,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2426-L2462","documentation":"In tcgen05.mma, the accumulator's leading dimension m must equal the LHS's leading dimension m (LHS is (m, k), accumulator is (m, n)). This check fires when acc.shape[0] != a.shape[0].","triggerScenarios":"Calling tcgen05.mma with an accumulator allocated as (m2, n) where m2 differs from the m rows of the LHS block, e.g. LHS (128, 16) with accumulator (64, 64).","commonSituations":"Mismatched BlockMapping causing the accumulator grid dimension to differ from the LHS block's M; manually allocated TMEM accumulator sized for a different tile shape.","solutions":["Size the accumulator's first dimension to match the LHS block's first dimension","Verify the BlockMapping/scratch shapes that allocate the TMEM accumulator","Re-derive accumulator shape as (a.shape[0], b.shape[1]) (or n*2 in collective mode)"],"exampleFix":"# before\nacc = allocate(TMEM, (64, n), jnp.float32)\ntcgen05.mma(a, b, acc, k_dim=16)\n# after\nacc = allocate(TMEM, (a.shape[0], n), jnp.float32)\ntcgen05.mma(a, b, acc, k_dim=16)","handlingStrategy":"validation","validationCode":"m, k = a.shape\nassert acc.shape[0] == m, f'acc m {acc.shape[0]} != lhs m {m}'\ntcgen05.mma(a, b, acc, k_dim=k)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive accumulator shape from operand shapes, never hardcode","Centralize tile-size constants"],"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"}