{"record":{"id":"e67d0fc503829ea2","repo":"jax-ml/jax","slug":"expected-an-accumulator-ref-got-acc","errorCode":null,"errorMessage":"Expected an accumulator ref, got {acc}","messagePattern":"Expected an accumulator ref, got (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/primitives.py","lineNumber":1293,"sourceCode":"      acc_transforms\n  )\n  matmul_acc_lhs_p.bind(\n      acc_ref,\n      lhs,\n      *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","sourceCodeStart":1275,"sourceCodeEnd":1311,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/primitives.py#L1275-L1311","documentation":"The matmul_acc_lhs primitive requires its first argument to be a Ref whose memory space is an AccMemorySpace (the MXU accumulator). If the ref lives in VMEM/SMEM/HBM, the abstract eval raises this ValueError.","triggerScenarios":"Calling matmul_acc_lhs on a ref allocated with MemorySpace.VMEM (or HBM/SMEM) instead of an accumulator-space ref; passing a plain array instead of a Ref.","commonSituations":"Allocating the output ref for the matmul without tpu_core.AccMemorySpace.ACC; refactoring a kernel that previously used VMEM output for matmul_pop.","solutions":["Allocate/write the accumulator into a ref with memory space tpu_core.MemorySpace.ACCUMULATOR (AccMemorySpace)","Pass the acc Ref (from kernel args) directly, not a transformed or copied block"],"exampleFix":"# before\nout = pallas_core.new_ref(..., memory_space=tpu_core.MemorySpace.VMEM)\nmatmul_acc_lhs(out, lhs)\n# after\nout = pallas_core.new_ref(..., memory_space=tpu_core.AccMemorySpace.ACCUMULATOR)\nmatmul_acc_lhs(out, lhs)","handlingStrategy":"type-guard","validationCode":"from jax._src.pallas.tpu import tpu_core\nassert isinstance(acc.memory_space, tpu_core.AccMemorySpace)","typeGuard":"def is_accumulator_ref(ref) -> bool:\n    return isinstance(getattr(ref, 'memory_space', None), tpu_core.AccMemorySpace)","tryCatchPattern":null,"preventionTips":["Allocate matmul accumulators in ACC memory space","Pass kernel-scope acc refs directly to matmul primitives"],"tags":["jax","pallas","tpu","matmul","memory-space"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}