jax-ml/jax · error · ValueError

Cannot swap scalars to VMEM.

Error message

Cannot swap scalars to VMEM.

What it means

Raised during a VMEM swap when the memory aval is vector-shaped but the output aval is scalar: swapping (read-modify-write returning old value) a scalar in VMEM is unsupported because the returned old value cannot be a bare vector-memory scalar.

Source

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

  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:
    result = tpu.strided_load(mem_aval_vec_type, ref, starts, strides)
  else:
    result = vector.load(mem_aval_vec_type, ref, starts)
  val = _maybe_cast_store_to_memref_type(ctx, val_aval, val)
  if mem_aval != aval_out:
    if not aval_out.shape:
      raise ValueError("Cannot swap scalars to VMEM.")
    # We are slicing a scalar so provided dummy 1 indices
    result_vec_type = ir.VectorType.get(
        ctx.lowering_context.dynamic_shape_replacement_fn(aval_out.shape),
      _dtype_to_ir_type(aval_out.dtype, is_kernel_boundary=True))
    result = vector.shape_cast(result_vec_type, result)
    val_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))
    val = vector.shape_cast(val_vec_type, val)
    if mask is not None:
      mask_vec_type = ir.VectorType.get(
          ctx.lowering_context.dynamic_shape_replacement_fn(mem_aval.shape),
          _dtype_to_ir_type(mask_aval.dtype)
      )
      mask = vector.shape_cast(mask_vec_type, mask)
  result = _maybe_cast_load_to_bool(ctx, val_aval, result)

  if need_stride:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Give the swapped ref shape (1,) so the old value is a 1-element vector
2. Avoid swap for scalars: use separate load then store

Example fix

# before
old = pl.swap(ref, scalar)  # ref block shape ()
# after
old = pl.swap(ref, scalar[None])[0]  # ref block shape (1,)

if val.shape == ():
    val, ref_block = val[None], (1,)  # swap with vector shape
old = pl.swap(ref_block_ref, val)[0]

When it happens

Trigger: pl.swap on a VMEM ref where the stored value/dtype physicalization makes mem_aval != aval_out and aval_out.shape is ().

Common situations: Using swap to atomically read-and-write a scalar accumulator in VMEM; extended dtypes (like custom element types) that force physicalized (vector) memory layouts.

Related errors

• {primitive_name}: Buffers with a memory space of HBM or ANY
• Cannot store scalars to VMEM
• Only arrays with 32-bit element types can be converted to sc
• Swap does not support storing to {ref_memory_space!r}. Copy
• Swap only supports slices with stride 1, got {strides}

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