jax-ml/jax · error · ValueError

Loads and stores are only allowed on VMEM and SMEM reference

Error message

Loads and stores are only allowed on VMEM and SMEM references.{extra}

What it means

Raised when a store targets a ref in a memory space other than VMEM or SMEM. Synchronous stores only work on vector and scalar memory; ANY-space buffers must be written via the async copy (DMA) path.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2525

  )
  need_stride = not all((s is None or s == 1) for s in strides)

  if is_smem_store:
    if mask is not None:
      raise ValueError("SMEM store does not support masks")
    if val_aval.shape:
      raise ValueError("Can only store scalars to SMEM")
    result = memref.load(ref, starts)
    result = _maybe_cast_load_to_bool(ctx, val_aval, result)
    val = _maybe_cast_store_to_memref_type(ctx, val_aval, val)
    memref.store(val, ref, starts)
    return result

  if not is_vmem_store:
    extra = ""
    if memory_space == "#tpu.memory_space<any>":
      extra = " ANY memory space can only be accessed using async_copy."
    raise ValueError(
        "Loads and stores are only allowed on VMEM and SMEM references." + extra
    )

  # handling VMEM store below
  if not val_aval.shape:
    raise ValueError("Cannot store scalars to VMEM")

  mem_slice_shape = list(aval_out.shape)
  for i, a in enumerate(idx.indices):
    if not isinstance(a, primitives.Slice):
      mem_slice_shape.insert(i, 1)
  mem_slice_shape_iter = iter(mem_slice_shape)
  mem_slice_shape = [
      1 if b is pallas_core.squeezed else next(mem_slice_shape_iter)
      for b in ref_block_shape
  ]
  mem_aval = aval_out.update(
      shape=tuple(mem_slice_shape), sharding=jax_core.get_cur_mesh_sharding()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Write into a VMEM scratch ref, then pltpu.async_copy(scratch, any_ref, ...) to move results out
2. Use pltpu.async_copy with the ANY-space buffer as the DMA source/destination only

Example fix

# before
pl.store(any_out_ref, val)
# after
pl.store(vmem_scratch, val)
pltpu.async_copy(vmem_scratch, any_out_ref)
pltpu.async_copy_wait()

def store_any_space(vmem_scratch, any_ref, val):
    pl.store(vmem_scratch, val)
    pltpu.async_copy(vmem_scratch, any_ref)
    pltpu.async_copy_wait()

def is_any_space(ref) -> bool:
    return 'any>' in str(getattr(ref.aval, 'memory_space', ''))

When it happens

Trigger: pl.store to a ref whose memory_space is ANY (or otherwise not vmem/smem), e.g. trying to write final results directly from an ANY-space output without async_copy.

Common situations: Mixing async_copy pipelines with direct stores; writing to the ANY-space copy destination/source directly.

Related errors

• Loads are only allowed on VMEM and SMEM references.
• Memory space {self.memory_space} is not supported by mesh {s
• Unsupported core type: {core_type}
• Invalid memory space: {memory_space!r}
• Unknown action: {action}

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