{"record":{"id":"cb95d85068643fb4","repo":"jax-ml/jax","slug":"async-store-only-supports-tmem-refs","errorCode":null,"errorMessage":"Async store only supports TMEM refs","messagePattern":"Async store only supports TMEM refs","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":4223,"sourceCode":"\n  Args:\n    ref: The TMEM reference to store to.\n    value: The value to store.\n  \"\"\"\n  ref, ref_transforms = state_primitives.get_ref_and_transforms(\n      ref, None, \"async_store_tmem\"\n  )\n  flat_ref_transforms, ref_transforms_treedef = tree_util.tree_flatten(\n      ref_transforms\n  )\n  async_store_tmem_p.bind(\n      ref, value, *flat_ref_transforms, tree=ref_transforms_treedef\n  )\n\n@async_store_tmem_p.def_effectful_abstract_eval\ndef _async_store_tmem_abstract_eval(ref, val, *avals_flat, tree):\n  if ref.memory_space != gpu_core.MemorySpace.TMEM:\n    raise ValueError(\"Async store only supports TMEM refs\")\n  _, effects = state_primitives._swap_abstract_eval(\n      ref, val, *avals_flat, tree=tree\n  )\n  return (), effects\n\n@lowering.register_lowering_rule(async_store_tmem_p, mgpu.LoweringSemantics.Lane)\ndef _async_store_tmem_lowering_rule(\n    ctx: lowering.LoweringRuleContext, x_ref, value, *leaves, tree\n):\n  assert isinstance(x_ref, tcgen05.TMEMRef)\n  x_aval = ctx.avals_in[0]\n  assert isinstance(x_aval, state_types.AbstractRef)\n  transforms = jax.tree.unflatten(tree, leaves)\n  transform_avals = tree.unflatten(\n      ctx.avals_in[2 : 2 + tree.num_leaves]\n  )\n  x_tmem, _, transforms = lowering._handle_transforms(\n      ctx, x_aval, x_ref, transform_avals, transforms, handle_transposes=False,","sourceCodeStart":4205,"sourceCodeEnd":4241,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L4205-L4241","documentation":"The async_store_tmem primitive only stores into Tensor Memory references. A destination ref whose memory_space is not TMEM raises ValueError in its abstract eval.","triggerScenarios":"Calling the async TMEM store API with an SMEM or HBM ref as destination.","commonSituations":"Storing results of a TMEM load back to the wrong buffer; copy-paste errors between load/store kernel templates.","solutions":["Store into a TMEM-allocated ref","Use regular async_copy or async_store for non-TMEM destinations"],"exampleFix":"// before\nasync_store_tmem(smem_ref, val)\n// after\nasync_store_tmem(tmem_ref, val)","handlingStrategy":"type-guard","validationCode":"assert ref.memory_space == gpu_core.MemorySpace.TMEM","typeGuard":"def is_tmem(ref): return getattr(ref, 'memory_space', None) == gpu_core.MemorySpace.TMEM","tryCatchPattern":null,"preventionTips":["Double-check destination memory space before async TMEM stores"],"tags":["jax","pallas","tmem","memory-space","blackwell"],"backgroundTag":"invalid-memory-space","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}