{"record":{"id":"51b084b2f3711042","repo":"jax-ml/jax","slug":"strided-slices-not-implemented","errorCode":null,"errorMessage":"Strided slices not implemented","messagePattern":"Strided slices not implemented","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":966,"sourceCode":"  if trailing_dims := len(shape) - len(index):\n    index += (slice(None),) * trailing_dims\n  base_indices: list[ir.Value | int] = []\n  slice_shape = []\n  is_squeezed = []\n  for axis, (idx, bound) in enumerate(zip(index, shape)):\n    if isinstance(idx, (ir.Operation, ir.OpView)):\n      idx = idx.result\n    if isinstance(idx, int):\n      if check_oob and (idx >= bound or (idx < 0 and -idx > bound)):\n        raise IndexError(\n            f\"Index {idx} along axis {axis} is out of bounds for shape {shape}\"\n        )\n      base_indices.append(idx if idx >= 0 else bound + idx)\n      slice_shape.append(1)\n      is_squeezed.append(True)\n    elif isinstance(idx, slice):\n      if idx.step is not None and idx.step != 1:\n        raise NotImplementedError(\"Strided slices not implemented\")\n      start = idx.start or 0\n      if start < 0:\n        start = bound + start\n      stop = idx.stop or bound\n      if stop < 0:\n        stop = bound + stop\n      if check_oob and (\n          start < 0 or start >= bound or stop < 0 or stop > bound\n      ):\n        raise IndexError(\n            f\"Slice {idx} along axis {axis} is out of bounds for shape {shape}\"\n        )\n      base_indices.append(start)\n      slice_shape.append(stop - start)\n      is_squeezed.append(False)\n    elif isinstance(idx, DynamicSlice):\n      if check_oob and (\n          isinstance(idx.base, int) and idx.base + idx.length > bound","sourceCodeStart":948,"sourceCodeEnd":984,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L948-L984","documentation":"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.","triggerScenarios":"ref[:, ::2], ref[1:5:2], or any slice whose .step is not None and != 1, passed to __getitem__, __setitem__, slice, or memref_slice.","commonSituations":"Porting numpy/jax.numpy slicing idioms into a Mosaic kernel; trying to downsample or reverse a shared-memory buffer with step slicing.","solutions":["Materialize the strided access yourself: loop over range(start, stop, step) and index element-by-element (guard OOB)","Copy the desired elements into a new memref with an explicit gather loop","Request/implement strided slice support upstream if the pattern is core to your kernel"],"exampleFix":"# before\nsub = buf[:, ::2]\n# after\nsub = memref.alloca(ir.MemRefType.get((rows, cols // 2), elem_ty), [], [])\nfor i in range(rows):\n  for j in range(0, cols, 2):\n    store(sub, [i, j // 2], load(buf, [i, j]))","handlingStrategy":"fallback","validationCode":"def safe_step(s: slice) -> bool:\n    return s.step is None or s.step == 1\nassert all(safe_step(i) for i in indices if isinstance(i, slice))","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid step-slicing in Mosaic kernels; write explicit loops with range(start, stop, step)","Reverse data by copying, not by [::-1]"],"tags":["jax","mosaic-gpu","memref","slice","not-implemented","strided"],"backgroundTag":"unsupported-slice-step","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}