jax-ml/jax · error · NotImplementedError

Swap only supports slices with stride 1, got {strides}

Error message

Swap only supports slices with stride 1, got {strides}

What it means

SC swap/store only supports 0-d (scalar) values in SMEM. Trying to swap an array in SMEM, or a scalar in a non-SMEM space, raises NotImplementedError with this message.

Source

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

  ref_block_shape, *_ = ctx.block_shapes
  ref, ref_block_shape = _transform_ref(
      ref, ref_aval, ref_block_shape, prev_transforms
  )
  starts, sizes, strides, squeeze_dims, _ = tc_lowering._indexer_to_start_size_stride(
      indexer, ref_block_shape, cast_to_index=True
  )
  for first_nontrivial_dim, s in enumerate(sizes):
    if s != 1:
      break
  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")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. For scalar swaps, move the scalar into SMEM first
2. For array swaps in SMEM, use element-wise stores instead of swap
3. Load/swap arrays in VMEM and copy scalars through SMEM

Example fix

# before
old = smem_buf[:]  # swap of array in SMEM
# after
old = smem_buf[0]  # scalar swap in SMEM only

assert (out.ndim == 0) == (space is MemorySpace.SMEM)

null

null

When it happens

Trigger: (out_aval.ndim == 0) mismatched with memory space: swapping an array-shaped value in SMEM, or a scalar in VMEM/HBM.

Common situations: Using SMEM buffers sized >1 element with swap semantics; scalar swaps attempted on VMEM refs.

Related errors

• Swap does not support storing to {ref_memory_space!r}. Copy
• Get does not support loading from {ref_memory_space!r}. Copy
• Swap only supports scalars in SMEM.
• Swap does not support masked scalar stores
• Gather only supports loading from VMEM, got {ref_aval.memory

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