jax-ml/jax · error · ValueError

Cannot store scalars to VMEM

Error message

Cannot store scalars to VMEM

What it means

Raised when storing a scalar (shape ()) value to VMEM. Vector memory ops on TPU are vector ops and cannot write a bare scalar; VMEM stores require at least 1-D values. Scalar stores belong in SMEM.

Source

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

    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()
  )
  mem_aval_vec_type = ir.VectorType.get(
      ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),
      _dtype_to_ir_type(mem_aval.dtype, is_kernel_boundary=True)
  )
  if need_stride:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape the scalar to (1,) and give the output block shape (1,): store val[None]
2. Move that output to SMEM and store the scalar there

Example fix

# before
pl.store(out_ref, scalar_val)  # out block shape () in VMEM
# after
pl.store(out_ref, scalar_val[None])  # out block shape (1,)

if val.shape == ():
    val = val[None]  # make it (1,) for VMEM store
pl.store(vmem_ref, val)

When it happens

Trigger: pl.store to a VMEM ref with val_aval.shape == ().

Common situations: Storing a per-block scalar (loss, norm) to an output with scalar block shape in VMEM.

Related errors

• {primitive_name}: Buffers with a memory space of HBM or ANY
• Cannot swap scalars to VMEM.
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}

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