jax-ml/jax · error · ValueError

Gather only supports loading from VMEM, got {ref_aval.memory

Error message

Gather only supports loading from VMEM, got {ref_aval.memory_space}

What it means

Raised by the SparseCore lowering rule for pallas_load/gather when the reference being loaded from is not in VMEM (Vector Memory). Mosaic's SparseCore gather only supports reading out of VMEM or the DEFAULT memory space; any other memory space (e.g. SMEM/CMEM) is rejected at lowering time.

Source

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

      raise ValueError(
          f"{mask.shape=} does not match the expected shape {out_aval.shape}"
      )
    if mask.dtype != jnp.bool:
      raise TypeError(f"Mask must be a boolean array, got {mask.dtype}")
  return out_aval, {state_types.ReadEffect(0)}


@sc_lowering.register_lowering_rule(gather_p)
def _gather_lowering_rule(
    ctx: sc_lowering.LoweringRuleContext, *flat_args, tree
):
  ref, transforms, indices, mask = tree.unflatten(flat_args)
  ref_aval, *_ = tree.unflatten(ctx.avals_in)
  if ref_aval.memory_space not in (
      tpu_core.MemorySpace.VMEM,
      pallas_core.MemorySpace.DEFAULT,
  ):
    raise ValueError(
        f"Gather only supports loading from VMEM, got {ref_aval.memory_space}"
    )
  if transforms:
    ref_block_shape, *_ = ctx.block_shapes
    ref, _ = tc_lowering._transform_ref(
        ref, ref_aval, ref_block_shape, transforms
    )
  [out_aval] = ctx.avals_out
  vec_type = ir.VectorType.get(
      out_aval.shape, sc_lowering._dtype_to_ir_type(ref_aval.dtype)
  )
  return tpu.vector_load_idx(vec_type, ref, indices, mask=mask)


def load_gather(
    ref: Ref, indices: Sequence[jax.Array], *, mask: jax.Array | None = None
) -> jax.Array:
  """Gathers an array from a ref.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Declare/annotate the gathered ref with memory space VMEM (tpu_core.MemorySpace.VMEM) or leave it DEFAULT
Check where the ref is created (kernel signature or scratch allocation) and fix its memory space
Avoid gather on non-VMEM buffers; copy to a VMEM buffer first inside the kernel

Example fix

// before
ref = pallas_core.new_memory_scope(...)  # memory_space=SMEM
val = sc_primitives.load(ref, ...)  # gather path

// after
# allocate/cast the ref in VMEM
val = kernel_call.load_from_vmem(ref, ...)

Defensive patterns

Strategy: validation

Validate before calling

ms = ref.aval.memory_space if hasattr(ref, 'aval') else ref.memory_space
assert ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT), ms

Type guard

def is_gatherable_ref(ref) -> bool:
    ms = getattr(getattr(ref, 'aval', ref), 'memory_space', None)
    return ms in (tpu_core.MemorySpace.VMEM, pallas_core.MemorySpace.DEFAULT)

Try / catch

try:
    v = load_scatterlike(ref, indices)
except ValueError as e:
    if 'only supports loading from VMEM' in str(e):
        raise RuntimeError(f'ref in unsupported memory space; move to VMEM: {e}')
    raise

Prevention

Declare all gathered refs in VMEM in kernel signatures
Write a helper that allocates scratch only in VMEM for SC kernels
Add unit asserts on memory spaces when building kernels for multiple core types

When it happens

Trigger: Calling a Pallas kernel using load/gather on a tpu_core.BlockLayout ref whose memory_space is not tpu_core.MemorySpace.VMEM or pallas_core.MemorySpace.DEFAULT, e.g. a ref annotated with a CoreMemorySpace pointing at a non-VMEM space.

Common situations: Writing a SparseCore (SC) kernel and declaring intermediate scratch buffers or refs in the wrong memory space; porting a TensorCore kernel to SparseCore where SMEM was used; upgrading JAX where memory-space checking became stricter.

Related errors

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