jax-ml/jax · error · ValueError

Can only load scalars from SMEM

Error message

Can only load scalars from SMEM

What it means

Raised when loading from an SMEM (scalar memory) ref with a non-scalar (array-shaped) output. SMEM on TPU is scalar memory; loads from it must produce scalars. Attempting to load a vector/array block from SMEM is a shape/memory-space mismatch.

Source

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

    physical_out_dtype = physical_element_aval.dtype
    physical_out_shape = jax_core.physical_shape(
        aval_out.shape, aval_out.dtype
    )
  else:
    physical_out_dtype = aval_out.dtype
    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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove SMEM memory_space from tensor inputs; use default VMEM for arrays
2. Reserve SMEM for scalars and PRNG key data only

# reject array-valued SMEM specs at setup time
for name, spec in specs.items():
    if 'SMEM' in str(getattr(spec, 'memory_space', '') or '') and any(d != 1 for d in (spec.block_shape or ())):
        raise ValueError(f'{name}: SMEM only supports scalar blocks')

When it happens

Trigger: pl.load where the ref's memory_space is SMEM but the output aval (or the requested slice) has shape with any non-1 dimensions.

Common situations: Setting memory_space=SMEM on a BlockSpec for a tensor (array) input to try to 'fix' another error; SMEM is only for scalars/PRNG keys.

Related errors

• Can only store scalars to SMEM
• PRNG keys must be loaded from SMEM. Did you set the memory s
• Indexing into a ()-shaped Ref not yet supported on TPU.
• Expected value and mask to have the same shape, but got valu
• SMEM store does not support masks

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