jax-ml/jax · error · NotImplementedError

Swap does not support masked scalar stores

Error message

Swap does not support masked scalar stores

What it means

Raised in the SparseCore store/swap lowering: masked stores are only supported for array accesses, so passing a mask to a scalar swap raises NotImplementedError. The input at fault is the mask argument of a scalar swap.

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
• Swap only supports scalars in SMEM.
• Swap only supports slices with stride 1, got
• Cannot swap scalars to VMEM.
• Gather only supports loading from VMEM, got

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

Appendix: 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)