{"record":{"id":"c0621221a138c6aa","repo":"jax-ml/jax","slug":"only-a-single-tiling-transform-is-supported-when-c","errorCode":null,"errorMessage":"Only a single tiling transform is supported when collapsing a shape, but got {in_transforms=} and {out_transforms=}","messagePattern":"Only a single tiling transform is supported when collapsing a shape, but got (.+?) and (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2265,"sourceCode":"# layout inference. It is not entirely clear what the best approach is.\ndef _check_collapse_shape(\n    op: memref.CollapseShapeOp,\n    in_transforms: Sequence[lc.MemRefTransform],\n    out_transforms: Sequence[lc.MemRefTransform],\n):\n  if len(in_transforms) != len(out_transforms):\n    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","sourceCodeStart":2247,"sourceCodeEnd":2283,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2247-L2283","documentation":"collapse_shape lowering supports at most a single TileTransform on input and output; anything else (multiple transforms, non-tiling transforms) is unimplemented.","triggerScenarios":"memref.collapse_shape with more than one input transform, or transforms that aren't lc.TileTransform on either side.","commonSituations":"Collapsing a tensor that has both tiling and swizzle transforms attached; hitting the single-transform limitation of the collapse path.","solutions":["Ensure exactly one TileTransform in and one out, nothing else","Strip swizzle transforms before the collapse (or use untiled memrefs)","Restructure the kernel so collapse happens on untransformed data"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def single_tile_transforms_only(in_ts, out_ts) -> bool:\n    return (len(in_ts) == 1 and len(out_ts) == 1\n            and isinstance(in_ts[0], lc.TileTransform)\n            and isinstance(out_ts[0], lc.TileTransform))","tryCatchPattern":null,"preventionTips":["Limit collapse inputs to one tiling transform","Strip swizzles before collapsing"],"tags":["jax","mosaic-gpu","collapse-shape","tile-transform","not-implemented"],"backgroundTag":"unsupported-layout-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}