{"record":{"id":"d8c3658c2e710c99","repo":"jax-ml/jax","slug":"collapseshapeop-with-non-contiguous-strides-is-not","errorCode":null,"errorMessage":"CollapseShapeOp with non-contiguous strides is not supported.","messagePattern":"CollapseShapeOp with non-contiguous strides is not supported\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2068,"sourceCode":") -> ConstraintSystemDerivationRuleResult:\n  reassociation = tuple(len(ir.ArrayAttr(idx)) for idx in op.reassociation)\n  # This should only occur when going from a (1, ...) shape to an empty shape.\n  # We can handle it if needed, but right now `CollapseShape` will not deal with\n  # this case.\n  if not reassociation:\n    raise NotImplementedError(\n        \"CollapseShapeOp with empty reassociation is not supported.\"\n    )\n\n  source = ValueSite(op, VariableType.OPERAND, 0)\n  source_var = ctx.producer_ref(source)\n  dest = ValueSite(op, VariableType.RESULT, 0)\n  dest_var = cs.Variable(dest)\n\n  strides, _ = ir.MemRefType(source.value.type).get_strides_and_offset()\n  # In this case, we'd need additional checks to produce a correct constraint.\n  if strides != utils.get_contiguous_strides(source.shape):\n    raise NotImplementedError(\n        \"CollapseShapeOp with non-contiguous strides is not supported.\"\n    )\n\n  # TODO(bchetioui): We could generate an inverse expression `ExpandShape` in\n  # order to allow inferring layouts bidirectionally. This would allow removing\n  # transforms from some kernels' BlockSpecs, but is not necessary at this time.\n  collapse_expr = cs.CollapseShape(source_var, source.shape, reassociation)\n  return cs.ConstraintSystem(constraints=[cs.Equals(dest_var, collapse_expr)]), {\n      source_var: [source],\n      dest_var: [dest],\n  }\n\n\n# `memref.load` and `memref.store` are used to load barrier phases which are\n# scalars---the rule needn't do anything interesting, but we need to have it.\n@_add_constraint_system_derivation_rule(memref.LoadOp)\n@_add_constraint_system_derivation_rule(memref.StoreOp)\ndef _memref_load_store_op_constraint_system(","sourceCodeStart":2050,"sourceCodeEnd":2086,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2050-L2086","documentation":"Mosaic GPU's layout inference handles memref.CollapseShapeOp only when the source memref has contiguous row-major strides. If the actual strides differ from get_contiguous_strides(shape) (e.g. after a transpose or strided subview), the pass raises NotImplementedError because producing correct layout constraints for non-contiguous collapse would require additional checks.","triggerScenarios":"Applying memref.collapse_shape to a memref with non-contiguous strides: a transposed memref, a strided subview (e.g. a column slice), or a swizzled smem layout, inside a Mosaic GPU kernel.","commonSituations":"Flattening a tile after slicing columns (row stride > row length), reshaping transposed shared-memory tiles, or collapsing views created by TMA/smem transforms with padded or swizzled strides.","solutions":["Materialize a contiguous copy of the memref (round-trip through registers or a fresh smem allocation) before collapse_shape","Reorder operations: collapse first on the contiguous source, then apply transpose/strided access on the collapsed result","If strides are contiguous but include a padding dimension, adjust the allocation/padding so get_strides_and_offset matches contiguous strides"],"exampleFix":"# before: src has strides [1, 512] but shape [512, 128] (padded/strided)\nflat = memref.collapse_shape(src, reassociation=[[0, 1]])\n\n# after: copy into contiguous buffer first\ncontig = alloc_contiguous_like(src)\ncopy(src, contig)\nflat = memref.collapse_shape(contig, reassociation=[[0, 1]])","handlingStrategy":"validation","validationCode":"strides, _ = ir.MemRefType(src.type).get_strides_and_offset()\nif strides != utils.get_contiguous_strides(src.shape):\n    src = copy_to_contiguous(src)  # materialize before collapse_shape","typeGuard":"def is_contiguous_memref(v) -> bool:\n    m = ir.MemRefType(v.type)\n    s, _ = m.get_strides_and_offset()\n    return s == utils.get_contiguous_strides(m.shape)","tryCatchPattern":null,"preventionTips":["Flatten/collapse only freshly allocated or copied buffers","Watch for padding/swizzle-induced stride mismatches on smem buffers before reshape"],"tags":["jax","mosaic-gpu","memref","collapse-shape","strides","layout-inference"],"backgroundTag":"unsupported-operation-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}