jax-ml/jax · error · NotImplementedError

Strided slices unsupported. Got stride: {ds.stride}

Error message

Strided slices unsupported. Got stride: {ds.stride}

What it means

TilingTransform.commute_ndindexer rewrites an indexer through a tiled view. Slicing is only supported with unit stride; a slice with stride None (dynamic) or any value != 1 raises NotImplementedError because strided access through tiled layouts is not implemented.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:817

    del aval
    idxs = indexer.indices
    indexer_shape = indexer.shape
    untiled_idxs = idxs[: -len(self.tiling)]
    tiled_idxs = idxs[-len(self.tiling) :]
    idxs_after_tiling: list[indexing.Slice] = []
    leading_shape, untiled_shape = (
        indexer_shape[: -len(self.tiling)],
        indexer_shape[-len(self.tiling) :],
    )
    for idx, tile, dim in zip(tiled_idxs, self.tiling, untiled_shape):
      match idx:
        case slice() | indexing.Slice():
          if isinstance(idx, slice):
            ds = indexing.Slice.from_slice(idx, dim)
          else:
            ds = idx
          if ds.stride is not None and ds.stride != 1:
            raise NotImplementedError(
                f"Strided slices unsupported. Got stride: {ds.stride}"
            )
          start, size = ds.start, ds.size
          if (
              start is not None and isinstance(start, int) and start % tile
          ) or (size is not None and isinstance(size, int) and size % tile):
            raise ValueError(
                f"Expected slice start ({start}) and slice size ({size})"
                f" to be divisible by the tile size ({tile})"
            )
          def _maybe_cdiv_with_cast(x, y):
            if x is None:
              return None
            if isinstance(x, jax.Array):
              # If x is an int32, we need to make sure y is an int32 to avoid
              # a dtype mismatch.
              y = jnp.array(y, x.dtype)
            return pallas_utils.cdiv(x, y)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace the strided slice with explicit indices (e.g. `ref[jnp.arange(0, n, 2)]` via arange indexing) if supported by the surrounding op
2. Materialize the strided read outside the tiled ref (copy to a normal array first)
3. Ensure slices have stride exactly 1 (or None-with-default-1 after fixing slice construction)

Example fix

# before
x = ref[:, ::2]  # NotImplementedError: strided slice through tiling

# after
idx = jnp.arange(0, ref.shape[1], 2)
x = ref[:, idx]

def slice_is_unit_stride(s):
    return getattr(s, 'stride', None) in (None, 1) or getattr(s, 'step', None) in (None, 1)

null

null

When it happens

Trigger: Indexing a tiled ref with a strided slice such as `ref[::2]` or `ref[0:128:4]`, or an indexing.Slice with dynamic stride, when the tiling transform must commute the indexer during lowering.

Common situations: Downsampling/strided access patterns in Pallas GPU kernels; converting existing indexing code to tiled block specs; passing slices with symbolic strides that default to non-1.

Related errors

• Get only supports slices with stride 1, got {strides}
• Expected slice start ({start}) and slice size ({size}) to be
• Transpose cannot be moved before a tiling transform when it
• Unsupported index type: {type(idx)}
• Commuting a `UntilingTransform` with a `ReshapeTransform` is

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