{"record":{"id":"437c7cfbe894cb55","repo":"jax-ml/jax","slug":"swap-does-not-support-storing-to-ref-memory-space","errorCode":null,"errorMessage":"Swap does not support storing to {ref_memory_space!r}. Copy the data to a core-local memory space, e.g. VMEM, via `pltpu.async_copy`.","messagePattern":"Swap does not support storing to (.+?)\\. Copy the data to a core-local memory space, e\\.g\\. VMEM, via `pltpu\\.async_copy`\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/sc_lowering.py","lineNumber":189,"sourceCode":"  )\n\n\ndef _store_lowering_rule(\n    ctx: LoweringRuleContext, ref, val, mask, *flat_transforms, tree, add\n):\n  ref_aval, _, *_flat_index_avals = ctx.avals_in\n  assert isinstance(ref_aval, state.AbstractRef)\n  [out_aval] = ctx.avals_out\n  assert isinstance(out_aval, jax_core.ShapedArray)\n\n  ref_memory_space = tpu_core.memory_space_to_tpu_memory_space(\n      ref_aval.memory_space, ctx.lowering_context.kernel_type\n  )\n  if (\n      ref_memory_space is MemorySpace.HBM\n      or ref_memory_space is MemorySpace.VMEM_SHARED\n  ):\n    raise NotImplementedError(\n        f\"Swap does not support storing to {ref_memory_space!r}.\"\n        \" Copy the data to a core-local memory space, e.g. VMEM,\"\n        \" via `pltpu.async_copy`.\"\n    )\n\n  transforms = list(tree_util.tree_unflatten(tree, flat_transforms))\n  if not transforms or not isinstance(transforms[-1], indexing.NDIndexer):\n    tref_aval = state.transform_type(transforms, ref_aval)\n    assert isinstance(tref_aval, state.AbstractRef)\n    transforms.append(indexing.NDIndexer.make_trivial_indexer(tref_aval.shape))\n  *prev_transforms, indexer = transforms\n  ref_block_shape, *_ = ctx.block_shapes\n  ref, ref_block_shape = _transform_ref(\n      ref, ref_aval, ref_block_shape, prev_transforms\n  )\n  starts, sizes, strides, squeeze_dims, _ = tc_lowering._indexer_to_start_size_stride(\n      indexer, ref_block_shape, cast_to_index=True\n  )","sourceCodeStart":171,"sourceCodeEnd":207,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/sc_lowering.py#L171-L207","documentation":"The Swap (ref write returning old value) lowering on SparseCore only supports core-local memory spaces. Swapping into HBM or VMEM_SHARED refs is unimplemented; stage the data via pltpu.async_copy instead.","triggerScenarios":"swap operations (e.g. pallas_swap / swap=True stores) targeting an HSM or VMEM_SHARED ref in an SC kernel.","commonSituations":"Accumulator swap patterns written for TensorCore kernels reused on SC; writing results directly to HBM refs with swap semantics.","solutions":["Swap into a VMEM buffer, then async_copy the buffer out to HBM","Use a plain store for HBM output if the old value is not needed"],"exampleFix":"# before\nold = hbm_ref[i, :]  # swap/store into HBM\n# after\npltpu.async_copy(vmem_buf, hbm_ref[i, :])\nold = vmem_buf[:]  # swap within VMEM only","handlingStrategy":"fallback","validationCode":"if ref_aval.memory_space in (MemorySpace.HBM, MemorySpace.VMEM_SHARED):\n    # route writes through async_copy instead","typeGuard":"null","tryCatchPattern":"null","preventionTips":["Swap only in VMEM/SMEM; flush to HBM via DMA"],"tags":["jax","pallas","tpu","sparsecore","memory-space","swap"],"backgroundTag":"unsupported-memory-space-access","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}