jax-ml/jax · error · ValueError

Expected an accumulator ref, got

Error message

Expected an accumulator ref, got {acc}

What it means

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.

Solutions

1. Allocate/write the accumulator into a ref with memory space tpu_core.MemorySpace.ACCUMULATOR (AccMemorySpace)
2. Pass the acc Ref (from kernel args) directly, not a transformed or copied block

Example fix

# before
out = pallas_core.new_ref(..., memory_space=tpu_core.MemorySpace.VMEM)
matmul_acc_lhs(out, lhs)
# after
out = pallas_core.new_ref(..., memory_space=tpu_core.AccMemorySpace.ACCUMULATOR)
matmul_acc_lhs(out, lhs)

from jax._src.pallas.tpu import tpu_core
assert isinstance(acc.memory_space, tpu_core.AccMemorySpace)

def is_accumulator_ref(ref) -> bool:
    return isinstance(getattr(ref, 'memory_space', None), tpu_core.AccMemorySpace)

When it happens

Trigger: 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.

Common situations: Allocating the output ref for the matmul without tpu_core.AccMemorySpace.ACC; refactoring a kernel that previously used VMEM output for matmul_pop.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• {ctx.avals_out[0].dtype}
• Gather only supports loading from VMEM, got
• Get does not support loading from
• Invalid memory space
• load_staged_rhs must be an integer or None.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/e67d0fc503829ea2. Report an issue: GitHub.

Appendix: source

Thrown at jax/_src/pallas/mosaic/primitives.py:1293

      acc_transforms
  )
  matmul_acc_lhs_p.bind(
      acc_ref,
      lhs,
      *flat_acc_transforms,
      acc_transforms_tree=acc_transforms_treedef,
      load_staged_rhs=load_staged_rhs,
  )


@matmul_acc_lhs_p.def_effectful_abstract_eval
def _matmul_acc_lhs_abstract_eval(
    acc: state.AbstractRef, lhs, *flat_acc_transforms, acc_transforms_tree, load_staged_rhs
):
  del load_staged_rhs,  # Unused.
  transforms = tree_util.tree_unflatten(acc_transforms_tree, flat_acc_transforms)
  if not isinstance(acc.memory_space, tpu_core.AccMemorySpace):
    raise ValueError(f"Expected an accumulator ref, got {acc}")
  transformed_acc = state.transform_type(transforms, acc)
  assert isinstance(transformed_acc, state.AbstractRef)
  acc_shape: tuple[int, ...] = transformed_acc.shape
  if len(acc_shape) != 2:
    raise ValueError(
        f"The shape of the accumulator {acc_shape} is not 2-dimensional."
    )
  m, _ = acc_shape
  if m != lhs.shape[0]:
    raise ValueError(
        f"The shape of the accumulator {acc_shape} does not "
        f"match the shape of the lhs {lhs.shape}."
    )
  return [], {mxu_effect, state.ReadEffect(0), state.WriteEffect(0)}


matmul_pop_p = jax_core.Primitive("matmul_pop")

View on GitHub (pinned to 1e1c6a8fc0)