jax-ml/jax · error · ValueError

Can only store scalars to SMEM

Error message

Can only store scalars to SMEM

What it means

Raised when storing a non-scalar (array-shaped) value to SMEM. TPU scalar memory only holds scalars; vector data must live in VMEM. Storing an array to an SMEM ref is a memory-space/shape mismatch.

Source

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

  (aval_out,) = ctx.avals_out
  if not isinstance(val, ir.Value):
    val = ir_constant(val, mlir_type=_dtype_to_ir_type(val_aval.dtype))
  if not is_smem_store and not ref_block_shape:
    raise NotImplementedError(
        "Indexing into a ()-shaped Ref not yet supported on TPU.")

  starts, _, 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_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")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use VMEM (default) for array-valued refs
2. Keep SMEM only for scalars and PRNG key material

if 'SMEM' in str(spec.memory_space or '') and any(d != 1 for d in (spec.block_shape or ())):
    raise ValueError('SMEM stores must be scalar; use VMEM for arrays')

When it happens

Trigger: pl.store where the ref memory space is SMEM but val_aval.shape is non-empty.

Common situations: Wrongly marking a tensor output as SMEM while copying an SMEM scalar-store pattern; only PRNG keys and scalars belong in SMEM.

Related errors

• Can only load scalars from 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/bad5e659facdb1cb. Report an issue: GitHub.