{"record":{"id":"33b82b734b85aadd","repo":"jax-ml/jax","slug":"expected-wgmmaabstractaccumulatorref-got-acc-aval","errorCode":null,"errorMessage":"Expected WGMMAAbstractAccumulatorRef got {acc_aval}","messagePattern":"Expected WGMMAAbstractAccumulatorRef got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":1903,"sourceCode":"\n  wgmma_ref_p.bind(\n      acc,\n      a,\n      b,\n      *acc_transforms_leaves,\n      *a_transforms_leaves,\n      *b_transforms_leaves,\n      acc_transforms_tree=acc_transforms_tree,\n      a_transforms_tree=a_transforms_tree,\n      b_transforms_tree=b_transforms_tree,\n  )\n\n\n@wgmma_ref_p.def_effectful_abstract_eval\ndef _wgmma_ref_effectful_abstract_eval(acc_aval, a_aval, b_aval, *_, **params):\n  del b_aval, params\n  if not isinstance(acc_aval, gpu_core.WGMMAAbstractAccumulatorRef):\n    raise TypeError(f\"Expected WGMMAAbstractAccumulatorRef got {acc_aval}\")\n  return (), {\n      gpu_core._wgmma_pipeline_effect,\n      state.WriteEffect(0),\n      state.ReadEffect(0),\n      state.ReadEffect(2),\n      *([state.ReadEffect(1)] if isinstance(a_aval, state.AbstractRef) else [])\n  }\n\n\ndef _wgmma_ref_pp_eqn(\n    eqn: jax_core.JaxprEqn,\n    context: jax_core.JaxprPpContext,\n    settings: jax_core.JaxprPpSettings,\n):\n  del settings\n  acc, a, b, *leaves = eqn.invars\n  transform_treedefs = [\n      eqn.params[\"acc_transforms_tree\"],","sourceCodeStart":1885,"sourceCodeEnd":1921,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L1885-L1921","documentation":"The wgmma_ref primitive (updating an accumulator in place) requires its first operand to be a WGMMAAbstractAccumulatorRef — the special accumulator SMEM buffer allocated via the dedicated allocator. Passing any other ref aval is a TypeError.","triggerScenarios":"Binding wgmma_ref (or the wgmma_increments style API) with an accumulator that is a normal SMEM/TrackedRef instead of a WGMMA accumulator ref.","commonSituations":"Allocating the accumulator with the generic SMEM allocator instead of the WGMMA accumulator allocator; refactoring kernels where acc refs were swapped with ordinary buffers.","solutions":["Allocate the accumulator using the WGMMA accumulator allocation API (mgpu.SMEM_ALLOCATOR / wgmma accumulator helper) so its aval is WGMMAAbstractAccumulatorRef.","Check that the accumulator ref is passed as the first positional argument to wgmma."],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":null,"typeGuard":"from jax._src.pallas.mosaic_gpu import gpu_core\ndef is_wgmma_acc(ref):\n    return isinstance(getattr(ref, 'aval', ref), gpu_core.WGMMAAbstractAccumulatorRef)","tryCatchPattern":null,"preventionTips":["Allocate accumulators with the dedicated WGMMA accumulator API.","Pass the accumulator as the first operand to wgmma."],"tags":["mosaic-gpu","pallas","wgmma","accumulator","type-error"],"backgroundTag":"invalid-operand-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}