jax-ml/jax · error · NotImplementedError

Non-unit strides not implemented.

Error message

Non-unit strides not implemented.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:4531

            mgpu.c(int(x), i32), (), is_signed=False
        )
      raise NotImplementedError(x)

    num_skipped = 0
    for i in range(len(current_indices)):
      # Integer indexers remove dimensions which should be
      # skipped by following indexers.
      if i in removed_dimensions:
        num_skipped += 1
        continue
      dim_indexer = indexer.indices[i - num_skipped]
      current_index = current_indices[i]
      assert isinstance(current_index, indexing.Slice)

      current_start_index = _ensure_idx_fa(current_index.start)
      if isinstance(dim_indexer, indexing.Slice):
        if dim_indexer.stride != 1:
          raise NotImplementedError("Non-unit strides not implemented.")
        current_indices[i] = indexing.Slice(
            current_start_index + _ensure_idx_fa(dim_indexer.start),
            dim_indexer.size,
            1,
        )
      else:
        current_indices[i] = current_start_index + _ensure_idx_fa(dim_indexer)
        removed_dimensions.add(i)
  return indexing.NDIndexer(
      indices=tuple(current_indices),
      shape=root_shape,
      int_indexer_shape=(),
  )


@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(primitives.semaphore_read_p, mgpu.LoweringSemantics.Warpgroup)
def _semaphore_read_lowering_rule(ctx: LoweringRuleContext, *args, args_tree):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace strided slices with explicit gather using iota + multiplication, or reshape trick
2. Materialize the strided array before passing it into the kernel
3. Track upstream support for strided slices in Mosaic GPU

Example fix

# before
x = ref[::2]
# after
idx = pl.program_id(0) * 2
x = pl.load(ref, (idx,))  # explicit strided access

When it happens

Trigger: 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.

Common situations: Porting code that strides over arrays (downsampling, gathering every k-th element) to pallas GPU kernels.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Slicing batch dimensions is not supported.
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

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