jax-ml/jax · error · NotImplementedError

Swap does not support storing to {ref_memory_space!r}. Copy

Error message

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`.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/sc_lowering.py:189

  )


def _store_lowering_rule(
    ctx: LoweringRuleContext, ref, val, mask, *flat_transforms, tree, add
):
  ref_aval, _, *_flat_index_avals = ctx.avals_in
  assert isinstance(ref_aval, state.AbstractRef)
  [out_aval] = ctx.avals_out
  assert isinstance(out_aval, jax_core.ShapedArray)

  ref_memory_space = tpu_core.memory_space_to_tpu_memory_space(
      ref_aval.memory_space, ctx.lowering_context.kernel_type
  )
  if (
      ref_memory_space is MemorySpace.HBM
      or ref_memory_space is MemorySpace.VMEM_SHARED
  ):
    raise NotImplementedError(
        f"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`."
    )

  transforms = list(tree_util.tree_unflatten(tree, flat_transforms))
  if not transforms or not isinstance(transforms[-1], indexing.NDIndexer):
    tref_aval = state.transform_type(transforms, ref_aval)
    assert isinstance(tref_aval, state.AbstractRef)
    transforms.append(indexing.NDIndexer.make_trivial_indexer(tref_aval.shape))
  *prev_transforms, indexer = transforms
  ref_block_shape, *_ = ctx.block_shapes
  ref, ref_block_shape = _transform_ref(
      ref, ref_aval, ref_block_shape, prev_transforms
  )
  starts, sizes, strides, squeeze_dims, _ = tc_lowering._indexer_to_start_size_stride(
      indexer, ref_block_shape, cast_to_index=True
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Swap into a VMEM buffer, then async_copy the buffer out to HBM
2. Use a plain store for HBM output if the old value is not needed

Example fix

# before
old = hbm_ref[i, :]  # swap/store into HBM
# after
pltpu.async_copy(vmem_buf, hbm_ref[i, :])
old = vmem_buf[:]  # swap within VMEM only

if ref_aval.memory_space in (MemorySpace.HBM, MemorySpace.VMEM_SHARED):
    # route writes through async_copy instead

null

null

When it happens

Trigger: swap operations (e.g. pallas_swap / swap=True stores) targeting an HSM or VMEM_SHARED ref in an SC kernel.

Common situations: Accumulator swap patterns written for TensorCore kernels reused on SC; writing results directly to HBM refs with swap semantics.

Related errors

• Swap only supports slices with stride 1, got {strides}
• Get does not support loading from {ref_memory_space!r}. Copy
• Swap only supports scalars in SMEM.
• Swap does not support masked scalar stores
• Gather only supports loading from VMEM, got {ref_aval.memory

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/437c7cfbe894cb55. Report an issue: GitHub.