{"record":{"id":"a900a70cf2831f6f","repo":"jax-ml/jax","slug":"incompatible-shapes-for-matrix-multiplication-lhs","errorCode":null,"errorMessage":"Incompatible shapes for matrix multiplication: lhs={a.shape}, rhs={b.shape=}, acc={acc.shape}","messagePattern":"Incompatible shapes for matrix multiplication: lhs=(.+?), rhs=(.+?), acc=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1852,"sourceCode":"  Conceptually, this is equivalent to doing ``acc[...] += a[...] @ b[...]``,\n  except that the computation is performed asynchronously.\n\n  Args:\n    acc: The accumulator reference. Needs to be allocated via\n      :func:`jax.experimental.pallas.run_scoped` called with a\n      :func:`jax.experimental.pallas.mosaic_gpu.WGMMAAccumulatorRef`.\n    a: The left hand side operand reference.\n    b: The right hand side operand reference.\n\n  See also:\n    :func:`jax.experimental.pallas.mosaic_gpu.wgmma_wait`\n  \"\"\"\n  m, n = acc.shape\n  m2, k = a.shape\n  k2, n2 = b.shape\n\n  if m != m2 or n != n2 or k != k2:\n    raise ValueError(\n        f\"Incompatible shapes for matrix multiplication: lhs={a.shape},\"\n        f\" rhs={b.shape=}, acc={acc.shape}\"\n    )\n\n  # A and B must share a dtype, except that the e4m3/e5m2 FP8 pair may be mixed:\n  # `wgmma` takes independent `.atype`/`.btype` operands for FP8.\n  fp8_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2)\n  both_fp8 = a.dtype in fp8_dtypes and b.dtype in fp8_dtypes\n  if a.dtype != b.dtype and not both_fp8:\n    raise ValueError(\n        \"Mixed input dtypes for matrix multiplication unsupported: \"\n        f\"lhs={a.dtype}, rhs={b.dtype}\"\n    )\n\n  acc_transforms_leaves: list\n  if isinstance(acc, pallas_core.TransformedRef):\n    acc_transforms_leaves, acc_transforms_tree = jax.tree.flatten(acc.transforms)\n    acc = acc.ref","sourceCodeStart":1834,"sourceCodeEnd":1870,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1834-L1870","documentation":"wgmma validates that the accumulator shape (m, n) and lhs (m, k), rhs (k, n) form a consistent matrix multiplication before emitting the instruction. Mismatched dimensions are a pure user API error.","triggerScenarios":"Calling wgmma(a, b, acc) where acc.shape != (a.shape[0], b.shape[1]) or a.shape[1] != b.shape[0].","commonSituations":"Swapped operands (b, a instead of a, b), wrong accumulator allocation after changing block sizes, or forgetting that wgmma expects (m,k)@(k,n).","solutions":["Check shapes: allocate acc as jnp.zeros((a.shape[0], b.shape[1]), ...) and ensure a.shape[1] == b.shape[0].","Verify operand order — lhs must be (m, k), rhs (k, n).","Add an assert in the kernel to fail fast with clearer context."],"exampleFix":"# before\nacc = smem.zeros((n_dim, m_dim), jnp.float32)\nout = wgmma(a, b, acc)\n# after\nacc = smem.zeros((a.shape[0], b.shape[1]), jnp.float32)\nout = wgmma(a, b, acc)","handlingStrategy":"validation","validationCode":"m, k = a.shape; k2, n = b.shape\nassert k == k2, f'inner dims {k} vs {k2}'\nassert acc.shape == (m, n), f'acc {acc.shape} != {(m, n)}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Allocate the accumulator from lhs/rhs shapes in the same expression as the wgmma call.","Add shape asserts in kernel setup code."],"tags":["mosaic-gpu","pallas","wgmma","shape-mismatch","matrix-multiply"],"backgroundTag":"matrix-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}