jax-ml/jax · error · ValueError

Loading from an accumulator is not supported. Use `matmul_po

Error message

Loading from an accumulator is not supported. Use `matmul_pop` instead, which will additionally zero out the accumulator.

What it means

Raised when a Pallas TPU kernel tries to load (read) from an accumulator memory space (AccMemorySpace, the MMA accumulator). Accumulators on TPU are write-pop semanantic; reading them requires matmul_pop, which also zeroes the accumulator. Direct loads are unsupported to prevent silently reading partially-consumed accumulator state.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2285

      if any(
          (not isinstance(a, primitives.Slice) and a.shape)
          for a in idx_aval.indices
      ):
        raise ValueError("Cannot do int indexing on TPU")
  return prev_transforms, idx


@register_lowering_rule(primitives.load_p, ensure_mlir_values=False)
def _load_lowering_rule(ctx: LoweringRuleContext, *args_flat, args_tree, **_):
  ref, transforms, mask, _ = args_tree.unflatten(args_flat)
  ref_aval, transforms_avals, _, _ = args_tree.unflatten(ctx.avals_in)
  prev_transforms, idx = _canonicalize_transforms_to_indexer(
      ref_aval, transforms, transforms_avals
  )
  if mask is not None:
    raise NotImplementedError
  if isinstance(ref_aval.memory_space, tpu_core.AccMemorySpace):
    raise ValueError(
        "Loading from an accumulator is not supported. Use `matmul_pop` "
        "instead, which will additionally zero out the accumulator."
    )

  ref_block_shape, *_ = ctx.block_shapes
  ref, ref_block_shape = _transform_ref(
      ref, ref_aval, ref_block_shape, prev_transforms
  )
  ref_type = ir.MemRefType(ref.type)
  is_smem_load = str(ref_type.memory_space) == "#tpu.memory_space<smem>"
  (aval_out,) = ctx.avals_out
  if isinstance(aval_out.dtype, prng.KeyTy) and pl_random.is_pallas_impl(
      aval_out.dtype._impl
  ):
    # TODO(justinfu): Merge this with standard extended dtype handling.
    if not is_smem_load:
      raise ValueError("PRNG keys must be loaded from SMEM. Did you set "
                       "the memory space to MemorySpace.SMEM in the "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the accumulator API: acc = pltpu.AMATRIX[...] style access or matmul_pop / accumulator .pop() to read and zero
2. If you only need the result, call .pop() on the accumulator which performs the read-and-zero
3. Restructure to store results into a VMEM ref instead of reading the accumulator

Example fix

// before
out = pl.load(acc_ref)  # acc_ref has memory_space AMEM
// after
out = acc_ref.pop()  # matmul_pop: reads and zeroes the accumulator

from jax._src.pallas.triton.tpu_core import AccMemorySpace  # conceptual
def assert_not_acc(ref_aval):
    if type(ref_aval.memory_space).__name__ == 'AccMemorySpace':
        raise ValueError('accumulator ref: use .pop()/matmul_pop, not load')

def is_accumulator_ref(ref) -> bool:
    return 'AccMemorySpace' in type(getattr(ref.aval, 'memory_space', None)).__name__

When it happens

Trigger: Calling pl.load / ref[...] on a Ref whose BlockSpec memory_space is tpu_core.AMEM (accumulator) inside a TPU Pallas kernel.

Common situations: Using the new accumulator matmul API (pltpu.MMA) and trying to read the accumulator ref with a normal load instead of accumulator.pop(...); migrating code from older accumulate-and-read patterns.

Related errors

• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}
• Accumulators are not available on TPU {info.chip_version}
• mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}
• The minor dimension size of an accumulator ref must be {info

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