{"record":{"id":"4daa941254c00303","repo":"jax-ml/jax","slug":"collapsing-the-shape-of-a-memref-with-non-contiguo","errorCode":null,"errorMessage":"Collapsing the shape of a memref with non-contiguous strides is not supported","messagePattern":"Collapsing the shape of a memref with non-contiguous strides is not supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2272,"sourceCode":"    raise ValueError(\n        \"Expected the same number of in/out transforms, but got \"\n        f\"{in_transforms=} and {out_transforms=}\"\n    )\n  if not in_transforms:\n    return\n  t_in, *in_transforms = in_transforms\n  t_out, *_ = out_transforms\n  if (in_transforms or\n      not isinstance(t_in, lc.TileTransform) or\n      not isinstance(t_out, lc.TileTransform)):\n    raise NotImplementedError(\n        \"Only a single tiling transform is supported when collapsing a shape, \"\n        f\"but got {in_transforms=} and {out_transforms=}\"\n    )\n  src_ty = ir.MemRefType(op.src.type)\n  strides, _ = src_ty.get_strides_and_offset()\n  if strides != utils.get_contiguous_strides(src_ty.shape):\n    raise NotImplementedError(\n        \"Collapsing the shape of a memref with non-contiguous strides is not \"\n        \"supported\"\n    )\n  reassociation = tuple(len(ir.ArrayAttr(idx)) for idx in op.reassociation)\n\n  collapsed_tiling = cs.reduce_expression(\n      cs.CollapseShape(cs.SMEMTransforms(t_in, None), tuple(src_ty.shape),\n                       reassociation),\n      {},\n  )\n\n  if isinstance(collapsed_tiling, cs.Unsatisfiable):\n    raise ValueError(f\"Input tiling {t_in.tiling} is not compatible with {op}\")\n\n  assert isinstance(collapsed_tiling, cs.SMEMTransforms)\n  expected_t_out = collapsed_tiling.tiling\n  assert expected_t_out is not None\n  if expected_t_out != t_out:","sourceCodeStart":2254,"sourceCodeEnd":2290,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2254-L2290","documentation":"Collapsing a tiled memref is only implemented when the source memref is fully contiguous. Non-contiguous strides (e.g. from a prior subview or slice) make the tiling reduction invalid.","triggerScenarios":"memref.collapse_shape on a tiled memref whose src_ty.get_strides_and_offset() != contiguous strides for its shape, e.g. collapsing a sliced or padded buffer.","commonSituations":"Chaining slice/subview (which produces strided views) followed by collapse_shape on smem tiles; or collapsing a padded allocation.","solutions":["Copy the strided data into a fresh contiguous buffer before collapsing","Collapse before taking subviews/slices","Ensure the allocation has no padding so strides are contiguous"],"exampleFix":"// before\ncollapsed = t.memref.collapse_shape(sliced_strided_ref, reassoc)\n// after\ncontig = t.copy_to_contiguous(sliced_strided_ref)  # or store/load\ncollapsed = t.memref.collapse_shape(contig, reassoc)","handlingStrategy":"validation","validationCode":"strides, _ = ir.MemRefType(src.type).get_strides_and_offset()\nassert strides == utils.get_contiguous_strides(ir.MemRefType(src.type).shape)","typeGuard":"def is_contiguous(ty) -> bool:\n    strides, _ = ty.get_strides_and_offset()\n    return strides == utils.get_contiguous_strides(ty.shape)","tryCatchPattern":null,"preventionTips":["Collapse before slicing","Materialize contiguous copies of strided views"],"tags":["jax","mosaic-gpu","collapse-shape","strides","contiguity"],"backgroundTag":"non-contiguous-strides","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}