jax-ml/jax · error · NotImplementedError

Strided slices not implemented

Error message

Strided slices not implemented

What it means

Mosaic's memref slicing (parse_indices) only supports slice objects with step 1 or None. Slices like a[::2] or a[::-1] require strided/reversed memory access that this MLIR lowering path does not implement, so it raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:966

  if trailing_dims := len(shape) - len(index):
    index += (slice(None),) * trailing_dims
  base_indices: list[ir.Value | int] = []
  slice_shape = []
  is_squeezed = []
  for axis, (idx, bound) in enumerate(zip(index, shape)):
    if isinstance(idx, (ir.Operation, ir.OpView)):
      idx = idx.result
    if isinstance(idx, int):
      if check_oob and (idx >= bound or (idx < 0 and -idx > bound)):
        raise IndexError(
            f"Index {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(idx if idx >= 0 else bound + idx)
      slice_shape.append(1)
      is_squeezed.append(True)
    elif isinstance(idx, slice):
      if idx.step is not None and idx.step != 1:
        raise NotImplementedError("Strided slices not implemented")
      start = idx.start or 0
      if start < 0:
        start = bound + start
      stop = idx.stop or bound
      if stop < 0:
        stop = bound + stop
      if check_oob and (
          start < 0 or start >= bound or stop < 0 or stop > bound
      ):
        raise IndexError(
            f"Slice {idx} along axis {axis} is out of bounds for shape {shape}"
        )
      base_indices.append(start)
      slice_shape.append(stop - start)
      is_squeezed.append(False)
    elif isinstance(idx, DynamicSlice):
      if check_oob and (
          isinstance(idx.base, int) and idx.base + idx.length > bound

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Materialize the strided access yourself: loop over range(start, stop, step) and index element-by-element (guard OOB)
2. Copy the desired elements into a new memref with an explicit gather loop
3. Request/implement strided slice support upstream if the pattern is core to your kernel

Example fix

# before
sub = buf[:, ::2]
# after
sub = memref.alloca(ir.MemRefType.get((rows, cols // 2), elem_ty), [], [])
for i in range(rows):
  for j in range(0, cols, 2):
    store(sub, [i, j // 2], load(buf, [i, j]))

def safe_step(s: slice) -> bool:
    return s.step is None or s.step == 1
assert all(safe_step(i) for i in indices if isinstance(i, slice))

When it happens

Trigger: ref[:, ::2], ref[1:5:2], or any slice whose .step is not None and != 1, passed to __getitem__, __setitem__, slice, or memref_slice.

Common situations: Porting numpy/jax.numpy slicing idioms into a Mosaic kernel; trying to downsample or reverse a shared-memory buffer with step slicing.

Related errors

• Only unit strides are supported but got {op.static_strides}.
• CollapseShapeOp with empty reassociation is not supported.
• Slice {idx} along axis {axis} is out of bounds for shape {sh
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.

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