jax-ml/jax · error · NotImplementedError

Integer indexing of refs that follows a non-trivial slice is

Error message

Integer indexing of refs that follows a non-trivial slice is not supported on SC

What it means

The SC Get lowering walks the ref index expression to find the first non-trivial (size != 1) dimension; any integer-index squeeze after that point cannot be lowered on SparseCore. Only leading trivial dims may be followed by scalar indexing.

Source

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

  if not transforms or not isinstance(transforms[-1], indexing.NDIndexer):
    tref_aval = state.transform_type(transforms, ref_aval)
    assert isinstance(tref_aval, state.AbstractRef)
    transforms.append(indexing.NDIndexer.make_trivial_indexer(tref_aval.shape))
  *prev_transforms, indexer = transforms
  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move integer indices before the non-trivial slice, or index the resulting array after the load: v = ref[slice]; v[i]
2. Reshape the ref so the integer-indexed dims come first

Example fix

# before
v = ref[pl.ds(start, n), j]
# after
v = ref[pl.ds(start, n)]
jth = v[j]

# ensure no squeeze dims after first nontrivial slice dim
assert not any(squeeze_dims[first_nontrivial_dim:])

null

null

When it happens

Trigger: Indexing like ref[ds(...), i] — a real slice along one dimension followed by an integer index along a later dimension — in a SparseCore kernel read.

Common situations: Writing ref[block_slice, scalar_idx] (valid on other backends) when porting kernels to SC.

Related errors

• Cannot do int indexing on TPU
• Can only load a single key per load.
• Dimension must be 0 for 1D iota.
• The current TPU does not have SparseCores
• Mesh has {self.num_cores} cores, but the current TPU chip ha

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