jax-ml/jax · error · NotImplementedError

Get only supports slices with stride 1, got

Error message

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

What it means

Raised in the SparseCore load lowering: ref gets are lowered to strided loads that only support unit stride, so a slice with stride != 1 raises NotImplementedError. The input at fault is the step of a pl.ds slice used to index a ref.

Solutions

  1. Load the contiguous range and decimate afterwards on the returned array
  2. Adjust data layout so the needed elements are contiguous

Example fix

# before
v = ref[0:128:2]
# after
v = ref[0:128]
v = v[::2]
Defensive patterns

Strategy: validation

Validate before calling

assert all(s == 1 for s in strides)

Type guard

null

Try / catch

null

Prevention

When it happens

Trigger: Using a step != 1 slice on a ref inside an SC kernel, e.g. ref[0:n:2] or ref[:, ::4].

Common situations: Downsampling/decimation patterns ported from TensorCore or CPU code.

Related errors

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

Appendix: source

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

  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(
        "Get only supports slices with stride 1, got {strides}"
    )

  if (out_aval.ndim == 0) != (ref_memory_space is MemorySpace.SMEM):
    message = "Get only supports loading scalars from SMEM."
    if ref_memory_space is MemorySpace.SMEM:
      message += " Trying to load an array of shape {out_aval.shape}."
    elif ref_memory_space is MemorySpace.VMEM:
      message += (
          " To load a scalar from VMEM, load an array first and then extract a"
          " particular element, e.g. ``v = ref[pl.ds(idx, ...)]; v[0]``."
      )
    else:
      message += f" Trying to load a scalar from {ref_memory_space!r}."
    raise NotImplementedError(message)
  if out_aval.ndim == 0:
    if mask is not None:
      raise NotImplementedError("Get does not support masked scalar loads")

View on GitHub (pinned to 1e1c6a8fc0)