{"record":{"id":"cfc65e23a45b9f84","repo":"jax-ml/jax","slug":"loading-from-an-accumulator-is-not-supported-use","errorCode":null,"errorMessage":"Loading from an accumulator is not supported. Use `matmul_pop` instead, which will additionally zero out the accumulator.","messagePattern":"Loading from an accumulator is not supported\\. Use `matmul_pop` instead, which will additionally zero out the accumulator\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2285,"sourceCode":"      if any(\n          (not isinstance(a, primitives.Slice) and a.shape)\n          for a in idx_aval.indices\n      ):\n        raise ValueError(\"Cannot do int indexing on TPU\")\n  return prev_transforms, idx\n\n\n@register_lowering_rule(primitives.load_p, ensure_mlir_values=False)\ndef _load_lowering_rule(ctx: LoweringRuleContext, *args_flat, args_tree, **_):\n  ref, transforms, mask, _ = args_tree.unflatten(args_flat)\n  ref_aval, transforms_avals, _, _ = args_tree.unflatten(ctx.avals_in)\n  prev_transforms, idx = _canonicalize_transforms_to_indexer(\n      ref_aval, transforms, transforms_avals\n  )\n  if mask is not None:\n    raise NotImplementedError\n  if isinstance(ref_aval.memory_space, tpu_core.AccMemorySpace):\n    raise ValueError(\n        \"Loading from an accumulator is not supported. Use `matmul_pop` \"\n        \"instead, which will additionally zero out the accumulator.\"\n    )\n\n  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n  ref_type = ir.MemRefType(ref.type)\n  is_smem_load = str(ref_type.memory_space) == \"#tpu.memory_space<smem>\"\n  (aval_out,) = ctx.avals_out\n  if isinstance(aval_out.dtype, prng.KeyTy) and pl_random.is_pallas_impl(\n      aval_out.dtype._impl\n  ):\n    # TODO(justinfu): Merge this with standard extended dtype handling.\n    if not is_smem_load:\n      raise ValueError(\"PRNG keys must be loaded from SMEM. Did you set \"\n                       \"the memory space to MemorySpace.SMEM in the \"","sourceCodeStart":2267,"sourceCodeEnd":2303,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2267-L2303","documentation":"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.","triggerScenarios":"Calling pl.load / ref[...] on a Ref whose BlockSpec memory_space is tpu_core.AMEM (accumulator) inside a TPU Pallas kernel.","commonSituations":"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.","solutions":["Use the accumulator API: acc = pltpu.AMATRIX[...] style access or matmul_pop / accumulator .pop() to read and zero","If you only need the result, call .pop() on the accumulator which performs the read-and-zero","Restructure to store results into a VMEM ref instead of reading the accumulator"],"exampleFix":"// before\nout = pl.load(acc_ref)  # acc_ref has memory_space AMEM\n// after\nout = acc_ref.pop()  # matmul_pop: reads and zeroes the accumulator","handlingStrategy":"validation","validationCode":"from jax._src.pallas.triton.tpu_core import AccMemorySpace  # conceptual\ndef assert_not_acc(ref_aval):\n    if type(ref_aval.memory_space).__name__ == 'AccMemorySpace':\n        raise ValueError('accumulator ref: use .pop()/matmul_pop, not load')","typeGuard":"def is_accumulator_ref(ref) -> bool:\n    return 'AccMemorySpace' in type(getattr(ref.aval, 'memory_space', None)).__name__","tryCatchPattern":null,"preventionTips":["Treat accumulator refs as write/pop-only in code review","Wrap accumulator reads behind a helper that calls pop()"],"tags":["jax","pallas","tpu","accumulator","matmul"],"backgroundTag":"tpu-accumulator-read","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}