{"record":{"id":"8080535875106438","repo":"jax-ml/jax","slug":"the-stored-value-has-shape-src-shape-but-the-ta","errorCode":null,"errorMessage":"The stored value has shape {src.shape}, but the target reference has shape {transformed_ref.shape}","messagePattern":"The stored value has shape (.+?), but the target reference has shape (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":575,"sourceCode":"    barrier_transforms_treedef,\n    **_,\n):\n  del cluster_idx  # Unused.\n  _check_ref(ref, \"ref\", gpu_core.SMEM)\n  _check_ref(barrier, \"barrier\", gpu_core.SMEM)\n  flat_ref_transforms_avals, flat_barrier_transforms_avals = util.split_list(\n      flat_transforms_avals,\n      [ref_transforms_treedef.num_leaves],\n  )\n  ref_transform_avals = ref_transforms_treedef.unflatten(\n      flat_ref_transforms_avals\n  )\n  barrier_transform_avals = barrier_transforms_treedef.unflatten(\n      flat_barrier_transforms_avals\n  )\n  transformed_ref = pallas_core.TransformedRef(ref, ref_transform_avals)\n  if src.shape != transformed_ref.shape:\n    raise TypeError(\n        f\"The stored value has shape {src.shape}, but the target reference has\"\n        f\" shape {transformed_ref.shape}\"\n    )\n  if src.dtype != transformed_ref.dtype:\n    raise TypeError(\n        f\"The stored value has dtype {src.dtype}, but the target reference has\"\n        f\" dtype {transformed_ref.dtype}\"\n    )\n  transformed_barrier = pallas_core.TransformedRef(barrier, barrier_transform_avals)\n  if transformed_barrier.size != 1:\n    raise TypeError(\n        \"Expected a single barrier, got a barrier reference with shape\"\n        f\" {transformed_barrier.shape}\"\n    )\n\n  effs = {gpu_core._memory_effect, state.WriteEffect(1)}\n  return (), effs\n","sourceCodeStart":557,"sourceCodeEnd":593,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L557-L593","documentation":"Abstract eval of async_store_smem compares the shape of the value being stored (src.shape) with the shape of the target SMEM reference after applying its transforms (TransformedRef). If they differ, a TypeError is raised because the async hardware store copies raw bits and cannot reshape.","triggerScenarios":"Calling async_store_smem(smem_ref, value, ...) where value.shape != smem_ref.shape after slicing/tiling transforms; e.g. storing a (128,) vector into a (16,8) tiled ref, or omitting a slice that reshapes the target.","commonSituations":"Off-by-one block sizes between compute and buffers; forgetting to apply the same slice to the ref as used for the source; changing block shapes without resizing SMEM buffers.","solutions":["Make src.shape exactly equal transformed_ref.shape (adjust block_size or reshape the value first)","Apply matching index/slice transforms to the SMEM ref so its transformed shape equals the value's shape","Store via a regular store (smem[...] = value) if reshaping semantics are needed"],"exampleFix":"# before\nasync_store_smem(smem, x, barrier)  # x.shape=(128,), smem.shape=(16,8)\n# after\nasync_store_smem(smem, x.reshape(smem.shape), barrier)","handlingStrategy":"validation","validationCode":"assert src.shape == tuple(t_ref.shape), (\n  f'async_store_smem shape mismatch: {src.shape} vs {t_ref.shape}')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Single-source block shapes from one constant used for both compute and buffers","Add shape asserts in kernel prologues during development"],"tags":["jax","pallas","mosaic-gpu","shape-mismatch","async-store"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}