jax-ml/jax · error · ValueError

Loads are only allowed on VMEM and SMEM references.

Error message

Loads are only allowed on VMEM and SMEM references.

What it means

Raised when a Pallas load targets a ref in a memory space other than VMEM or SMEM. Only vector memory (VMEM) and scalar memory (SMEM) support synchronous loads; other spaces (notably ANY) require the asynchronous DMA copy path (async_copy).

Source

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

    physical_out_shape = aval_out.shape
  if not is_smem_load and not ref_block_shape:
    raise NotImplementedError(
        "Indexing into a ()-shaped Ref not yet supported on TPU.")
  starts, sizes, strides, _, _ = _indexer_to_start_size_stride(
      idx,
      ref_block_shape,
      cast_to_index=True,
  )
  need_stride = not all((s is None or s == 1) for s in strides)
  if is_smem_load:
    if ctx.avals_out[0].shape:
      raise ValueError("Can only load scalars from SMEM")
    return _maybe_cast_load_to_bool(ctx, aval_out, memref.load(ref, starts))
  elif str(ref_type.memory_space) != "#tpu.memory_space<vmem>":
    extra = ""
    if str(ref_type.memory_space) == "#tpu.memory_space<any>":
      extra = " ANY memory space can only be accessed using async_copy."
    raise ValueError(
        "Loads are only allowed on VMEM and SMEM references." + extra
    )
  load_aval = jax_core.ShapedArray(sizes, dtype=physical_out_dtype)
  if need_stride:
    load_val = tpu.strided_load(
        ctx.aval_to_ir_type(load_aval, is_kernel_boundary=True),
        ref,
        starts,
        strides,
    )
  else:
    load_val = vector.load(
        ctx.aval_to_ir_type(load_aval, is_kernel_boundary=True),
        ref,
        starts,
    )
  if load_aval != aval_out:
    if physical_out_shape:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use async_copy to move data from ANY space into a VMEM scratch ref, then load from that ref
2. Or specify memory_space=VMEM explicitly in the BlockSpec if synchronous load is intended

Example fix

# before
x = pl.load(any_ref)  # any_ref has MemorySpace.ANY
# after
buf = pltpu.make_partitioned(...)  # VMEM scratch
pltpu.async_copy(any_ref, buf, ...)
pltpu.async_copy_wait(...)
x = pl.load(buf)

def load_any_space(ref, vmem_scratch):
    if str(getattr(ref.aval, 'memory_space', 'vmem')).endswith('any>'):
        pltpu.async_copy(ref, vmem_scratch); pltpu.async_copy_wait()
        return pl.load(vmem_scratch)
    return pl.load(ref)

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

When it happens

Trigger: pl.load on a ref in e.g. '#tpu.memory_space<any>' or HBM-typed space instead of VMEM; using a BlockSpec that left memory space as ANY and reading it directly.

Common situations: Using pltpu.async_copy for DMA but also trying to pl.load the same ANY-space ref; forgetting to copy from ANY space into VMEM before compute.

Related errors

• Loads and stores are only allowed on VMEM and SMEM reference
• 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/97d91a6ca410f657. Report an issue: GitHub.