jax-ml/jax · error · ValueError

Scatter only supports storing to VMEM, got {memory_space}

Error message

Scatter only supports storing to VMEM, got {memory_space}

What it means

The SparseCore scatter lowering only stores into VMEM (or DEFAULT). Storing directly to SMEM/CMEM/other memory spaces raises this ValueError during lowering.

Source

Thrown at jax/_src/pallas/mosaic/sc_primitives.py:365

def _scatter_lowering_rule(
    ctx: sc_lowering.LoweringRuleContext, *flat_args, tree, add
):
  ref, transforms, indices, x, mask = jax.tree.unflatten(tree, flat_args)
  ref_aval, *_ = tree.unflatten(ctx.avals_in)
  if isinstance(ref_aval.memory_space, pallas_core.CoreMemorySpace):
    if not isinstance(ref_aval.memory_space.mesh, sc_core.VectorSubcoreMesh):
      raise ValueError(
          "Scatter only supports VectorSubcoreMesh, got"
          f" {type(ref_aval.memory_space.mesh)}"
      )
    memory_space = ref_aval.memory_space.memory_space
  else:
    memory_space = ref_aval.memory_space
  if memory_space not in (
      tpu_core.MemorySpace.VMEM,
      pallas_core.MemorySpace.DEFAULT,
  ):
    raise ValueError(
        f"Scatter only supports storing to VMEM, got {memory_space}"
    )
  if transforms:
    ref_block_shape, *_ = ctx.block_shapes
    ref, _ = tc_lowering._transform_ref(
        ref, ref_aval, ref_block_shape, transforms
    )
  tpu.vector_store_idx(x, ref, indices, mask=mask, add=add)
  return ()


def store_scatter(
    ref: Ref,
    indices: Sequence[jax.Array],
    x: jax.Array,
    *,
    mask: jax.Array | None = None,
) -> None:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Perform the scatter into a VMEM buffer, then copy to the target space
2. Change the ref's memory space to VMEM/DEFAULT
3. Restructure so only supported spaces are scatter targets

Example fix

// before
sc_primitives.store_scatter(smem_ref, idx, x)

// after
sc_primitives.store_scatter(vmem_ref, idx, x)
# then copy vmem_ref -> smem_ref if needed

ms = ref_aval.memory_space
if isinstance(ms, pallas_core.CoreMemorySpace):
    ms = ms.memory_space
assert ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT)

def is_storable_space(ref_aval) -> bool:
    ms = ref_aval.memory_space
    if isinstance(ms, pallas_core.CoreMemorySpace):
        ms = ms.memory_space
    return ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT)

When it happens

Trigger: Calling store_scatter/addupdate_scatter on a ref allocated in SMEM or another non-VMEM CoreMemorySpace.memory_space.

Common situations: Optimizing scratch buffers into SMEM for TC kernels and reusing that code on SC; writing output directly from a non-VMEM intermediate.

Related errors

• Get does not support loading from {ref_memory_space!r}. Copy
• 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
• {x.shape=} does not match expected shape {expected_shape}

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