jax-ml/jax · error · NotImplementedError

Swap does not support masked scalar stores

Error message

Swap does not support masked scalar stores

What it means

Atomic scalar add via swap semantics is not implemented in the SC lowering (the TODO notes memref.atomic_rmw is unsupported by the SC compiler). swap_add on scalars therefore raises.

Source

Thrown at jax/_src/pallas/mosaic/sc_lowering.py:233

        "Integer indexing of refs that follows a non-trivial slice is not"
        " supported on SC"
    )
  if not all(s == 1 for s in strides):
    raise NotImplementedError(
        "Swap only supports slices with stride 1, got {strides}"
    )

  if (out_aval.ndim == 0) != (ref_memory_space is MemorySpace.SMEM):
    message = "Swap only supports scalars in SMEM."
    if ref_memory_space is MemorySpace.SMEM:
      message += " Trying to swap an array of shape {out_aval.shape}."
    else:
      message += f" Trying to swap a scalar in {ref_memory_space!r}."
    raise NotImplementedError(message)

  if out_aval.ndim == 0:
    if mask is not None:
      raise NotImplementedError("Swap does not support masked scalar stores")
    if add:
      # TODO(slebedev): We can use memref.atomic_rmw here, but the SC compiler
      # doesn't support it yet.
      raise NotImplementedError("Swap does not support atomic scalar adds")
    old_val = memref.load(ref, starts)
    memref.store(val, ref, starts)
    return old_val

  if not ctx.lowering_context.needs_layout_passes:
    _check_aval_is_supported("Swap", out_aval)
  out_vec_type = ir.VectorType.get(
      out_aval.shape, _dtype_to_ir_type(out_aval.dtype)
  )
  if not ctx.lowering_context.needs_layout_passes:
    old_val = tpu.vector_load(out_vec_type, ref, starts, strides=[], mask=mask)
    _ = tpu.vector_store(
        val, ref, indices=starts, strides=[], mask=mask, add=add
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Accumulate into a vector element instead: load a 1-element vector, add, store back
2. Serialize the update (only one lane performs it) and use a normal load-add-store
3. Accumulate in VMEM arrays and reduce at the end

Example fix

# before
ref[i] += delta  # scalar atomic add on SC
# after
v = ref[pl.ds(i, 1)]
v = v.at[0].add(delta)
ref[pl.ds(i, 1)] = v

# accumulate via 1-element vector instead of scalar atomic add
v = ref[pl.ds(i, 1)]
ref[pl.ds(i, 1)] = v + delta

null

null

When it happens

Trigger: Using swap-add (e.g. ref[idx] += v via swap with add=True) where the value is a scalar in a SparseCore kernel.

Common situations: Counter increments or scalar accumulations written with atomic add semantics on SC.

Related errors

• Swap does not support storing to {ref_memory_space!r}. Copy
• Swap only supports slices with stride 1, got {strides}
• Swap only supports scalars in SMEM.
• Cannot swap scalars to VMEM.
• The current TPU does not have SparseCores

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