{"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":"Shape-consistency guard in the wgmma wrapper: the accumulator must be (m, n), the lhs operand (m, k), and the rhs operand (k, n). The error fires when any of these dimension pairs disagree, meaning the operands cannot form a valid blocked matrix multiply accumulating into the supplied acc ref.","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","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-23T08:17:48.524Z"}