jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

SparseCore loads only lower contiguous slices: every stride in the index expression must be 1. Strided slices like ref[::2] raise NotImplementedError in the Get lowering rule.

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)

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]

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

null

null

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

• The current TPU does not have SparseCores
• Mesh has {self.num_cores} cores, but the current TPU chip ha
• You can't use two different ScalarSubcoreMeshes.
• {self} should have the same core axis name and number of cor
• {self} should have a different axis name from the TensorCore

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