jax-ml/jax · error · NotImplementedError

`strides` must contain only 1s.

Error message

`strides` must contain only 1s.

What it means

Mosaic's lowering of vector.extract_strided_slice only handles unit strides (a contiguous slice); strided extraction is not implemented.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:770

def _vector_shape_cast_op_lowering_rule(
    _: LoweringContext, op: vector.ShapeCastOp
) -> Sequence[ir.Value]:
  [layout] = inference_utils.in_layouts(op)
  out_vec_ty = ir.VectorType(op.result.type)
  assert out_vec_ty.has_static_shape
  a = _fragmented_array_from_ir(op.source, layout)
  return [
      fragmented_array_to_ir(a.reshape(tuple(out_vec_ty.shape)), out_vec_ty)
  ]


@_register_lowering(vector.ExtractStridedSliceOp)
def _vector_extract_strided_slice_op_lowering_rule(
    ctx: LoweringContext, op: vector.ExtractStridedSliceOp
) -> Sequence[ir.Value]:
  del ctx
  if any(ir.IntegerAttr(s).value != 1 for s in op.strides):
    raise NotImplementedError("`strides` must contain only 1s.")
  [in_layout] = inference_utils.in_layouts(op)
  [out_layout] = inference_utils.out_layouts(op)
  assert in_layout == out_layout
  out_vec_ty = ir.VectorType(op.result.type)
  assert out_vec_ty.has_static_shape
  a = _fragmented_array_from_ir(op.source, in_layout)
  indices = tuple(
      utils.DynamicSlice(
          ir.IntegerAttr(offset).value, ir.IntegerAttr(length).value
      )
      for offset, length in zip(op.offsets, op.sizes, strict=True)
  )
  result = a[indices]
  assert result.layout == layouts_lib.from_layout_attr(out_layout)
  return [fragmented_array_to_ir(result, out_vec_ty)]


@_register_lowering(vector.ExtractOp)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use unit strides and slice a contiguous range
2. For strided access, do it via loads/stores to memory or reshape/split ops instead
3. Combine extract with layout casts to get the elements you need

Example fix

// before
vector.extract_strided_slice v[0:8:2]
// after
vector.extract_strided_slice v[0:8:1]

assert all(ir.IntegerAttr(s).value == 1 for s in op.strides), 'use unit strides'

When it happens

Trigger: Emitting vector.extract_strided_slice with any stride attribute != 1 inside a Mosaic kernel.

Common situations: Trying to subsample a vector (take every other element) rather than a contiguous slice.

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

• Only slicing with static indices allowed.
• Unsupported reduction kind: {op.kind}
• Strides {strides} have lower rank than tiling {tiling}
• Can not tile strides when tiled dimensions have been transpo
• Stride {s} is not divisible by {d} (tile size = {t}). Stride

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