{"record":{"id":"30aee7be28def626","repo":"jax-ml/jax","slug":"accumulator-dtype-inner-dtype-does-not-match-val","errorCode":null,"errorMessage":"Accumulator dtype {inner.dtype} does not match value dtype {val.dtype}","messagePattern":"Accumulator dtype (.+?) does not match value dtype (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2355,"sourceCode":"\n\n@wgmma_accumulator_store_p.def_effectful_abstract_eval\ndef _wgmma_accumulator_store_abstract_eval(acc, val):\n  # Before discharge acc is a WGMMAAbstractAccumulatorRef. After discharge,\n  # the discharge rule re-binds the primitive and acc becomes a ShapedArray.\n  if isinstance(acc, gpu_core.WGMMAAbstractAccumulatorRef):\n    inner = acc.inner_aval\n    assert isinstance(inner, jax_core.ShapedArray)\n  elif isinstance(acc, jax_core.ShapedArray):\n    inner = acc\n  else:\n    raise TypeError(f\"Expected WGMMAAbstractAccumulatorRef or ShapedArray, got {type(acc)}\")\n  if inner.shape != val.shape:\n    raise ValueError(\n        f\"Accumulator shape {inner.shape} does not match value shape {val.shape}\"\n    )\n  if inner.dtype != val.dtype:\n    raise ValueError(\n        f\"Accumulator dtype {inner.dtype} does not match value dtype {val.dtype}\"\n    )\n  effects: set[jax_core.Effect] = {gpu_core._wgmma_pipeline_effect}\n  if isinstance(acc, gpu_core.WGMMAAbstractAccumulatorRef):\n    effects.add(state.WriteEffect(0))\n  return inner, effects\n\n\n@discharge.register_discharge_rule(wgmma_accumulator_store_p)\ndef _wgmma_accumulator_store_discharge(ctx, acc, val):\n  del ctx\n  return (wgmma_accumulator_store_p.bind(acc, val), None), []\n\n\n@lowering.register_lowering_rule(\n    wgmma_accumulator_store_p, mgpu.LoweringSemantics.Lane\n)\ndef _wgmma_accumulator_store_lowering(","sourceCodeStart":2337,"sourceCodeEnd":2373,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2337-L2373","documentation":"The WGMMA accumulator store requires the stored value's dtype to exactly match the accumulator's inner dtype (typically fp32). Storing a half/bfloat16 result directly fails.","triggerScenarios":"Computing the update in f16/bf16 (e.g. output of an op that downcasts) and storing into an fp32 WGMMA accumulator, or converting the loaded value with .astype(jnp.float16) before storing back.","commonSituations":"Mixed-precision kernels where the matmul operands are fp8/fp16 but accumulators are fp32; forgetting to cast the delta back after arithmetic in lower precision.","solutions":["Cast the value to the accumulator dtype before storing: val.astype(acc.inner_aval.dtype)","Keep all arithmetic on the accumulator value in fp32","If intentional, make the cast explicit rather than relying on implicit promotion"],"exampleFix":"# before\nwgmma_accumulator_store(acc, (delta).astype(jnp.bfloat16))\n# after\nwgmma_accumulator_store(acc, delta.astype(acc.inner_aval.dtype))","handlingStrategy":"validation","validationCode":"val = val.astype(acc.inner_aval.dtype)\nwgmma_accumulator_store(acc, val)","typeGuard":"def dtype_ok(acc, val):\n    inner = getattr(acc, 'inner_aval', acc)\n    return inner.dtype == val.dtype","tryCatchPattern":null,"preventionTips":["Standardize on fp32 for accumulator arithmetic","Cast explicitly at store boundaries"],"tags":["jax","pallas","mosaic-gpu","wgmma","dtype-mismatch"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}