{"record":{"id":"3bbb6d7cbff9797e","repo":"jax-ml/jax","slug":"the-shape-of-the-accumulator-acc-shape-is-not-2","errorCode":null,"errorMessage":"The shape of the accumulator {acc_shape} is not 2-dimensional.","messagePattern":"The shape of the accumulator (.+?) is not 2-dimensional\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1298,"sourceCode":"      *flat_acc_transforms,\n      acc_transforms_tree=acc_transforms_treedef,\n      load_staged_rhs=load_staged_rhs,\n  )\n\n\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.","sourceCodeStart":1280,"sourceCodeEnd":1316,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1280-L1316","documentation":"The MXU accumulator used by matmul_acc_lhs must be a 2D [M, N] tile because the TPU MXU performs 2D matrix products. A 1D or 3D+ accumulator shape is rejected.","triggerScenarios":"Calling matmul_acc_lhs with an accumulator ref whose shape has rank != 2, e.g. shape (128,) or (8, 128, 128).","commonSituations":"Allocating the accumulator with the batch dimension left in; using a flat scratch buffer instead of an [M, N] tile.","solutions":["Reshape/allocate the accumulator as a 2D [M, N] tile matching the MXU dimensions (e.g. 128x128)","Handle batching with a loop over batch index, keeping each accumulator 2D"],"exampleFix":"# before\nacc = new_ref((batch, 128, 128), ..., memory_space=ACC)\n# after\nfor b in range(batch):\n  acc = new_ref((128, 128), ..., memory_space=ACC)\n  matmul_acc_lhs(acc, lhs[b])","handlingStrategy":"validation","validationCode":"assert len(acc_shape) == 2, f'acc must be 2D, got {acc_shape}'","typeGuard":"def is_2d_shape(shape) -> bool:\n    return len(tuple(shape)) == 2","tryCatchPattern":null,"preventionTips":["Keep MXU accumulators [M, N]; handle batching with loops"],"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"}