{"record":{"id":"d8f44fc93a462510","repo":"jax-ml/jax","slug":"only-unit-strides-are-supported-but-got-op-static","errorCode":null,"errorMessage":"Only unit strides are supported but got {op.static_strides}.","messagePattern":"Only unit strides are supported but got (.+?)\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":1937,"sourceCode":") -> ConstraintSystemDerivationRuleResult:\n  source = ValueSite(op, VariableType.OPERAND, 0)\n  var_source_dest = ctx.producer_ref(source)\n  dest = ValueSite(op, VariableType.RESULT, 0)\n  return cs.ConstraintSystem(), {var_source_dest: [source, dest]}\n\n\n@_add_constraint_system_derivation_rule(memref.SubViewOp)\ndef _memref_subview_constraint_system(\n    ctx: DerivationContext,\n    op: memref.SubViewOp,\n) -> ConstraintSystemDerivationRuleResult:\n  source = ValueSite(op, VariableType.OPERAND, 0)\n  source_var = ctx.producer_ref(source)\n  result = ValueSite(op, VariableType.RESULT, 0)\n  result_var = cs.Variable(result)\n\n  if any(s != 1 for s in op.static_strides):\n    raise NotImplementedError(\n        f\"Only unit strides are supported but got {op.static_strides}.\"\n    )\n\n  # Collect all the constraints from all dimensions.\n  tiling_multiple = []\n  dynamic_offset_index = 0\n  for i, size in enumerate(op.static_sizes):\n    offset = op.static_offsets[i]\n    if offset == ir.ShapedType.get_dynamic_size():\n      offset = op.offsets[dynamic_offset_index]\n      dynamic_offset_index += 1\n\n    # Drop all dimensions up to and including the last dynamic size. Dynamic\n    # sizes are not supported yet.\n    #\n    # Supporting dynamic sizes here can be done analogously to how dynamic\n    # offsets are supported. The reason we don't support dynamic sizes now is\n    # because the lowering does not yet support them.","sourceCodeStart":1919,"sourceCodeEnd":1955,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L1919-L1955","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before: strided subview\nsv = memref_subview(smem, offsets=[0,0], sizes=[128,64], strides=[2,1])\n\n# after: contiguous subview, stride in registers\nsv = memref_subview(smem, offsets=[0,0], sizes=[128,64], strides=[1,1])\n# then take every other row after loading into registers","handlingStrategy":"validation","validationCode":"# before emitting a subview-derived op:\nstrides = tuple(op.static_strides)  # or your known slice step\nif any(s != 1 for s in strides):\n    raise ValueError('rewrite slice: non-unit strides unsupported by Mosaic layout inference')","typeGuard":"def has_unit_strides(static_strides: tuple[int, ...]) -> bool:\n    return all(s == 1 for s in static_strides)","tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","mosaic-gpu","memref","layout-inference","strides","not-implemented"],"backgroundTag":"unsupported-operation-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}