{"record":{"id":"c7cab9e50b5b7c57","repo":"jax-ml/jax","slug":"non-unit-strides-not-implemented","errorCode":null,"errorMessage":"Non-unit strides not implemented.","messagePattern":"Non-unit strides not implemented\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":4531,"sourceCode":"            mgpu.c(int(x), i32), (), is_signed=False\n        )\n      raise NotImplementedError(x)\n\n    num_skipped = 0\n    for i in range(len(current_indices)):\n      # Integer indexers remove dimensions which should be\n      # skipped by following indexers.\n      if i in removed_dimensions:\n        num_skipped += 1\n        continue\n      dim_indexer = indexer.indices[i - num_skipped]\n      current_index = current_indices[i]\n      assert isinstance(current_index, indexing.Slice)\n\n      current_start_index = _ensure_idx_fa(current_index.start)\n      if isinstance(dim_indexer, indexing.Slice):\n        if dim_indexer.stride != 1:\n          raise NotImplementedError(\"Non-unit strides not implemented.\")\n        current_indices[i] = indexing.Slice(\n            current_start_index + _ensure_idx_fa(dim_indexer.start),\n            dim_indexer.size,\n            1,\n        )\n      else:\n        current_indices[i] = current_start_index + _ensure_idx_fa(dim_indexer)\n        removed_dimensions.add(i)\n  return indexing.NDIndexer(\n      indices=tuple(current_indices),\n      shape=root_shape,\n      int_indexer_shape=(),\n  )\n\n\n@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Lane)\n@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Warpgroup)\ndef _semaphore_read_lowering_rule(ctx: LoweringRuleContext, *args, args_tree):","sourceCodeStart":4513,"sourceCodeEnd":4549,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L4513-L4549","documentation":"When merging chained slice indexers, the Mosaic GPU backend only supports slices with stride 1. If any indexing.Slice in the merge chain has a non-unit stride, this NotImplementedError fires.","triggerScenarios":"Indexing a reference in a Mosaic GPU kernel with a strided slice such as ref[::2] or ref[1:10:3], especially after chained indexing.","commonSituations":"Porting code that strides over arrays (downsampling, gathering every k-th element) to pallas GPU kernels.","solutions":["Replace strided slices with explicit gather using iota + multiplication, or reshape trick","Materialize the strided array before passing it into the kernel","Track upstream support for strided slices in Mosaic GPU"],"exampleFix":"# before\nx = ref[::2]\n# after\nidx = pl.program_id(0) * 2\nx = pl.load(ref, (idx,))  # explicit strided access","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never use strided slices ref[::k] in Mosaic GPU kernels","Express strides via explicit index arithmetic"],"tags":["pallas","mosaic-gpu","slicing","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}