jax-ml/jax · error · NotImplementedError

Get does not support loading from {ref_memory_space!r}. Copy

Error message

Get does not support loading from {ref_memory_space!r}. Copy the data to a core-local memory space, e.g. VMEM, via `pltpu.async_copy`.

What it means

On SparseCore, the Get (ref read) lowering only reads core-local memory spaces. Loading directly from HBM or VMEM_SHARED refs is not implemented; data must first be moved with a DMA copy (pltpu.async_copy) into VMEM/SMEM.

Source

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

  return _load_lowering_rule(ctx, ref, None, *flat_transforms, tree=tree)


def _load_lowering_rule(
    ctx: LoweringRuleContext, ref, mask, *flat_transforms, tree
):
  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"Get does not support loading from {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. Add an async_copy (DMA) from HBM into a VMEM scratch buffer and read from that
2. Restructure the kernel so reads come from SMEM/VMEM refs only

Example fix

# before
v = hbm_ref[i, :]  # direct read from HBM in SC kernel
# after
pltpu.async_copy(hbm_ref[i, :], vmem_buffer)  # in copy pipeline stage
v = vmem_buffer[:]

from jax._src.pallas.mosaic.tpu_core import MemorySpace
if ref_aval.memory_space in (MemorySpace.HBM, MemorySpace.VMEM_SHARED):
    raise SystemExit('stage via pltpu.async_copy first')

null

null

When it happens

Trigger: Reading an HSM/VMEM_SHARED ref with Python getitem (ref[...]) inside a SparseCore Pallas kernel, i.e. calling pallas_get on an HBM-backed ref.

Common situations: Porting a TPU TensorCore kernel that directly indexes HBM refs; forgetting the async_copy pipeline stage in SC kernels.

Related errors

• Swap does not support storing to {ref_memory_space!r}. Copy
• Swap only supports slices with stride 1, got {strides}
• Gather only supports loading from VMEM, got {ref_aval.memory
• Memory space {self.memory_space} is not supported by mesh {s
• Unsupported core type: {core_type}

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