jax-ml/jax · error · NotImplementedError

Swap only supports scalars in SMEM.

Error message

Swap only supports scalars in SMEM.

What it means

Raised in the SparseCore store/swap lowering: scalar swaps are only allowed in SMEM; a scalar swap in another memory space (or an array swap in SMEM) raises this NotImplementedError. The input at fault is the memory space of the ref targeted by the swap.

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)
Defensive patterns

Strategy: fallback

Validate before calling

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

Type guard

null

Try / catch

null

Prevention

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

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

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