jax-ml/jax · error · NotImplementedError

Only unit strides are supported but got {op.static_strides}.

Error message

Only unit strides are supported but got {op.static_strides}.

What it means

Mosaic GPU's layout inference pass builds a constraint system for memref.SubViewOp and only supports unit strides. When the subview's static strides contain any value other than 1 (e.g. produced by slicing with a step), the derivation rule raises NotImplementedError because non-unit strides would require stride-aware layout constraints that are not implemented.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:1937

) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  var_source_dest = ctx.producer_ref(source)
  dest = ValueSite(op, VariableType.RESULT, 0)
  return cs.ConstraintSystem(), {var_source_dest: [source, dest]}


@_add_constraint_system_derivation_rule(memref.SubViewOp)
def _memref_subview_constraint_system(
    ctx: DerivationContext,
    op: memref.SubViewOp,
) -> ConstraintSystemDerivationRuleResult:
  source = ValueSite(op, VariableType.OPERAND, 0)
  source_var = ctx.producer_ref(source)
  result = ValueSite(op, VariableType.RESULT, 0)
  result_var = cs.Variable(result)

  if any(s != 1 for s in op.static_strides):
    raise NotImplementedError(
        f"Only unit strides are supported but got {op.static_strides}."
    )

  # Collect all the constraints from all dimensions.
  tiling_multiple = []
  dynamic_offset_index = 0
  for i, size in enumerate(op.static_sizes):
    offset = op.static_offsets[i]
    if offset == ir.ShapedType.get_dynamic_size():
      offset = op.offsets[dynamic_offset_index]
      dynamic_offset_index += 1

    # Drop all dimensions up to and including the last dynamic size. Dynamic
    # sizes are not supported yet.
    #
    # Supporting dynamic sizes here can be done analogously to how dynamic
    # offsets are supported. The reason we don't support dynamic sizes now is
    # because the lowering does not yet support them.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove the stride: slice into a contiguous subview and apply striding manually in register space (e.g. via reshape/slicing after the load) instead of in the memref subview
Check op.static_strides of the generated memref.subview (dump the MLIR module with MAINT_RENDER.../debug printing) and rewrite the slice producing non-unit strides
If the stride is logically 1 but encoded differently, normalize the view so MLIR constant-folds strides to 1 before layout inference

Example fix

# before: strided subview
sv = memref_subview(smem, offsets=[0,0], sizes=[128,64], strides=[2,1])

# after: contiguous subview, stride in registers
sv = memref_subview(smem, offsets=[0,0], sizes=[128,64], strides=[1,1])
# then take every other row after loading into registers

Defensive patterns

Strategy: validation

Validate before calling

# before emitting a subview-derived op:
strides = tuple(op.static_strides)  # or your known slice step
if any(s != 1 for s in strides):
    raise ValueError('rewrite slice: non-unit strides unsupported by Mosaic layout inference')

Type guard

def has_unit_strides(static_strides: tuple[int, ...]) -> bool:
    return all(s == 1 for s in static_strides)

Prevention

Avoid stepped slices (x[::k]) on memrefs inside Mosaic kernels; stride in registers instead
Dump the MLIR module and assert all memref.subview ops have unit static_strides before layout inference

When it happens

Trigger: Calling tiling/slicing helpers that lower to memref.subview with a non-unit step (e.g. x[::2] style slices or strided views on smem/gmem memrefs) inside a Mosaic GPU kernel; any op.static_strides tuple with an entry != 1.

Common situations: Kernels that stride over tiles (e.g. processing every other row/column), or transforms in BlockSpecs that introduce strided subviews of shared memory. Often appears after upgrading Mosaic when previously-passing strided patterns now hit layout inference.

Related errors

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