jax-ml/jax · error · NotImplementedError

Swap only supports scalars in SMEM.

Error message

Swap only supports scalars in SMEM.

What it means

When swapping a scalar on SC, a mask cannot be applied — the lowering raises immediately for masked scalar stores. Bounds must be handled by clamping or branching instead.

Source

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

  else:
    first_nontrivial_dim = len(sizes)
  if any(squeeze_dims[first_nontrivial_dim:]):
    raise NotImplementedError(
        "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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Clamp or guard the index so no mask is needed
2. Swap a 1-element vector with a mask and extract element 0 afterwards

Example fix

# before
old = pallas_swap(ref, idx, val, mask=m)  # scalar + mask
# after
idx = jnp.where(m, idx, 0)
old = pallas_swap(ref, idx, val)

if mask is not None and out.ndim == 0:
    idx = jnp.where(mask, idx, 0)

null

null

When it happens

Trigger: swap=True store with mask set and 0-d value in a SparseCore kernel.

Common situations: Conditional scatter-style updates using masked swaps.

Related errors

• Swap does not support storing to {ref_memory_space!r}. Copy
• Swap only supports slices with stride 1, got {strides}
• Swap does not support masked scalar stores
• 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/24150c2d5773912e. Report an issue: GitHub.