{"record":{"id":"deb05095ee4c1474","repo":"jax-ml/jax","slug":"the-shape-of-the-accumulator-acc-shape-does-not","errorCode":null,"errorMessage":"The shape of the accumulator {acc_shape} does not match the shape of the lhs {lhs.shape}.","messagePattern":"The shape of the accumulator (.+?) does not match the shape of the lhs (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1303,"sourceCode":"\n@matmul_acc_lhs_p.def_effectful_abstract_eval\ndef _matmul_acc_lhs_abstract_eval(\n    acc: state.AbstractRef, lhs, *flat_acc_transforms, acc_transforms_tree, load_staged_rhs\n):\n  del load_staged_rhs,  # Unused.\n  transforms = tree_util.tree_unflatten(acc_transforms_tree, flat_acc_transforms)\n  if not isinstance(acc.memory_space, tpu_core.AccMemorySpace):\n    raise ValueError(f\"Expected an accumulator ref, got {acc}\")\n  transformed_acc = state.transform_type(transforms, acc)\n  assert isinstance(transformed_acc, state.AbstractRef)\n  acc_shape: tuple[int, ...] = transformed_acc.shape\n  if len(acc_shape) != 2:\n    raise ValueError(\n        f\"The shape of the accumulator {acc_shape} is not 2-dimensional.\"\n    )\n  m, _ = acc_shape\n  if m != lhs.shape[0]:\n    raise ValueError(\n        f\"The shape of the accumulator {acc_shape} does not \"\n        f\"match the shape of the lhs {lhs.shape}.\"\n    )\n  return [], {mxu_effect, state.ReadEffect(0), state.WriteEffect(0)}\n\n\nmatmul_pop_p = jax_core.Primitive(\"matmul_pop\")\n\n\ndef matmul_pop(acc: Ref) -> jax.Array:\n  \"\"\"Returns the result of a matrix multiplication from a specific MXU and zeroes the accumulator.\n\n  If the result is not ready yet (the MXU is still busy), the operation blocks.\n\n  ```{warning}\n  The kernel must not leave any data in the accumulator upon exit.\n  ```\n","sourceCodeStart":1285,"sourceCodeEnd":1321,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1285-L1321","documentation":"In matmul_acc_lhs, the leading dimension M of the 2D accumulator must equal lhs.shape[0], because each matmul step accumulates lhs @ rhs into the acc rows. A mismatch is rejected at abstract-eval time.","triggerScenarios":"Calling matmul_acc_lhs(acc, lhs) where acc.shape[0] != lhs.shape[0], e.g. acc of shape (128, 512) with lhs of shape (64, 512).","commonSituations":"Tiling lhs with a block size that differs from the accumulator's M dimension; mixing tile sizes between the lhs pipeline and the acc allocation.","solutions":["Make lhs block M equal acc.shape[0] (both typically 128 on TPU MXU)","Recompute tiling so lhs blocks are [acc_m, k]"],"exampleFix":"# before\nmatmul_acc_lhs(acc, lhs)  # acc (128,512), lhs (64,512)\n# after\nlhs = lhs.reshape(2, 64, 512)  # loop over 64-row slabs, or reallocate acc with M=64","handlingStrategy":"validation","validationCode":"assert acc_shape[0] == lhs.shape[0], 'M dims must match'","typeGuard":"def matmul_shapes_compatible(acc_shape, lhs_shape) -> bool:\n    return len(acc_shape) == 2 and acc_shape[0] == lhs_shape[0]","tryCatchPattern":null,"preventionTips":["Derive tiling constants (e.g. 128) from one shared config"],"tags":["jax","pallas","tpu","matmul","shape-mismatch"],"backgroundTag":"shape-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}