{"record":{"id":"19d11807469d728a","repo":"jax-ml/jax","slug":"accumulator-shape-inner-shape-does-not-match-val","errorCode":null,"errorMessage":"Accumulator shape {inner.shape} does not match value shape {val.shape}","messagePattern":"Accumulator shape (.+?) does not match value shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":2351,"sourceCode":"def wgmma_accumulator_store(acc_ref, val):\n  if not isinstance(acc_ref.aval, gpu_core.WGMMAAbstractAccumulatorRef):\n    raise TypeError(f\"acc must be a WGMMAAccumulatorAbstractRef, got {acc_ref.aval=}\")\n  wgmma_accumulator_store_p.bind(acc_ref, val)\n\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","sourceCodeStart":2333,"sourceCodeEnd":2369,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L2333-L2369","documentation":"Raised by the abstract eval of wgmma_accumulator_store when the value being stored into a WGMMA accumulator ref has a different shape than the accumulator's inner ShapedArray. The store must write back a tensor of exactly the same shape that was read from the accumulator.","triggerScenarios":"Calling wgmma_accumulator_store(acc, val) where val was reshaped/broadcast after the corresponding wgmma_accumulator_load, e.g. acc has shape (128, 8) but the computed value is (64, 16) or has extra leading dims.","commonSituations":"Reshaping the accumulator value between load and store, using a different BlockMapping/grid that changes block shapes, or passing the result of a matmul whose N dimension differs from the accumulator's.","solutions":["Ensure val keeps the exact shape returned by wgmma_accumulator_load","Check any reshape/broadcast between accumulator load and store and remove or compensate for it","Print acc.inner_aval.shape and val.shape right before the store to find the divergence"],"exampleFix":"# before\nacc_val = wgmma_accumulator_load(acc)\nout = acc_val.reshape(other_shape)\nwgmma_accumulator_store(acc, out)\n# after\nacc_val = wgmma_accumulator_load(acc)\nout = acc_val + delta  # same shape\nwgmma_accumulator_store(acc, out)","handlingStrategy":"validation","validationCode":"assert val.shape == acc.inner_aval.shape, (val.shape, acc.inner_aval.shape)\nwgmma_accumulator_store(acc, val)","typeGuard":"def can_store(acc, val):\n    inner = getattr(acc, 'inner_aval', acc)\n    return getattr(inner, 'shape', None) == val.shape","tryCatchPattern":"try:\n    wgmma_accumulator_store(acc, val)\nexcept ValueError as e:\n    if 'does not match value shape' in str(e):\n        wgmma_accumulator_store(acc, val.reshape(inner.shape))\n    else:\n        raise","preventionTips":["Keep loaded accumulator values unreshaped until after the store","Add shape asserts right after wgmma_accumulator_load"],"tags":["jax","pallas","mosaic-gpu","wgmma","shape-mismatch"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}