jax-ml/jax · error · ValueError

SMEM store does not support masks

Error message

SMEM store does not support masks

What it means

Raised when a store to SMEM (scalar memory) supplies a mask. SMEM stores compile to scalar memref.store ops which have no masking support on TPU.

Source

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

  is_smem_store = memory_space == "#tpu.memory_space<smem>"
  is_vmem_store = memory_space == "#tpu.memory_space<vmem>"
  (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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the mask for SMEM stores; guard with control flow or clamp the index instead
2. If conditional scalar writes are needed, compute the value to write (e.g. select old vs new) and store unconditionally

Example fix

// before
pl.store(smem_ref, val, mask=cond)
// after
val = jnp.where(cond, val, pl.load(smem_ref))
pl.store(smem_ref, val)

def smem_store(ref, val, cond):
    if cond is None:
        pl.store(ref, val)
    else:
        pl.store(ref, jnp.where(cond, val, pl.load(ref)))

def is_smem_ref(ref) -> bool:
    return str(getattr(ref.aval, 'memory_space', '')).endswith('smem')

When it happens

Trigger: pl.store(ref, val, mask=m) where ref's memory space is SMEM.

Common situations: Reusing a generic masked-store helper for both VMEM tensors and SMEM scalars; storing a scalar with a scalar predicate mask.

Related errors

• PRNG keys must be loaded from SMEM. Did you set the memory s
• Indexing into a ()-shaped Ref not yet supported on TPU.
• Can only load scalars from SMEM
• masked swap with non-32-bit data
• Expected value and mask to have the same shape, but got valu

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