{"record":{"id":"a5dca158513be16d","repo":"jax-ml/jax","slug":"the-stored-value-has-dtype-src-dtype-but-the-ta","errorCode":null,"errorMessage":"The stored value has dtype {src.dtype}, but the target reference has dtype {transformed_ref.dtype}","messagePattern":"The stored value has dtype (.+?), but the target reference has dtype (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":580,"sourceCode":"  _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\n\n@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Lane)\n@lowering.register_lowering_rule(async_store_smem_p, mgpu.LoweringSemantics.Warpgroup)\ndef _async_store_smem_lowering(\n    ctx: lowering.LoweringRuleContext,","sourceCodeStart":562,"sourceCodeEnd":598,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L562-L598","documentation":"async_store_smem's abstract eval checks that the dtype of the stored value matches the dtype of the transformed target SMEM reference. Mismatch raises TypeError because the async copy is bit-exact with no implicit casting.","triggerScenarios":"Calling async_store_smem where src.dtype differs from the SMEM buffer's dtype, e.g. storing a float32 tensor into a bfloat16 SMEM block, or int32 into int8.","commonSituations":"Allocating SMEM scratch without an explicit dtype so it defaults differently from the computed value; mixing precision regimes (bf16 compute, fp32 buffers) in fused kernels; refactoring a kernel to half precision but keeping old scratch dtypes.","solutions":["Allocate the SMEM block with the same dtype as the value (e.g. pl.SMEM((128,), x.dtype))","Explicitly cast the value before storing: x.astype(smem.dtype)","If a wider accumulation type is intended, use two stores or a regular store with a cast"],"exampleFix":"# before\nsmem = pl.SMEM((128,), jnp.float32)\nasync_store_smem(smem, x_bf16, barrier)\n# after\nasync_store_smem(smem, x_bf16.astype(jnp.float32), barrier)","handlingStrategy":"validation","validationCode":"assert src.dtype == smem_ref.dtype, (\n  f'async_store_smem dtype mismatch: {src.dtype} vs {smem_ref.dtype}')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always allocate SMEM scratch with dtype=value.dtype","Centralize dtype decisions (precision config) in one place"],"tags":["jax","pallas","mosaic-gpu","dtype-mismatch","async-store"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}