{"record":{"id":"ab5b657320e321e3","repo":"jax-ml/jax","slug":"input-output-tiling-mismatch-when-attempting-to-co","errorCode":null,"errorMessage":"Input/output tiling mismatch when attempting to collapse a shape. Expected output tiling to be {expected_t_out.tiling} for input tiling {t_in.tiling}, but got {t_out.tiling}","messagePattern":"Input/output tiling mismatch when attempting to collapse a shape\\. Expected output tiling to be (.+?) for input tiling (.+?), but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2291,"sourceCode":"        \"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:\n    raise ValueError(\n        \"Input/output tiling mismatch when attempting to collapse a shape. \"\n        f\"Expected output tiling to be {expected_t_out.tiling} for input \"\n        f\"tiling {t_in.tiling}, but got {t_out.tiling}\"\n    )\n\n\n@_register_lowering(memref.CollapseShapeOp, support_warp_semantics=True)\ndef _memref_collapse_shape_op_lowering_rule(\n    ctx: LoweringContext, op: memref.CollapseShapeOp\n) -> Sequence[ir.Value]:\n  del ctx\n\n  [in_transforms_attr] = inference_utils.in_transforms(op)\n  [out_transforms_attr] = inference_utils.out_transforms(op)\n\n  in_swizzle = swizzle_from_transforms_attr(in_transforms_attr)\n  in_transforms = memref_transforms_from_transforms_attr(in_transforms_attr)\n  out_swizzle = swizzle_from_transforms_attr(out_transforms_attr)","sourceCodeStart":2273,"sourceCodeEnd":2309,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2273-L2309","documentation":"For a collapsible tiled memref, the output tiling is fully determined by the input tiling and the reassociation. This error fires when the annotated output tiling differs from the computed expected tiling.","triggerScenarios":"memref.collapse_shape with a manually annotated out transform whose tiling != collapsed_tiling.tiling derived from t_in and the reassociation.","commonSituations":"Hand-writing the output transform instead of letting inference derive it, causing a mismatch with the canonical collapsed tiling.","solutions":["Let inference compute the output tiling (omit the manual out transform)","Or set the out tiling to the expected collapsed tiling reported in the message","Recompute: run cs.reduce_expression(cs.CollapseShape(...)) mentally/externally to get the right value"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"expected = cs.reduce_expression(cs.CollapseShape(cs.SMEMTransforms(t_in, None), tuple(shape), reassoc), {})\nassert not isinstance(expected, cs.Unsatisfiable) and expected.tiling.tiling == t_out.tiling","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't hand-annotate output tiling; use inference","Derive out tiling via the symbolic collapse"],"tags":["jax","mosaic-gpu","collapse-shape","tiling","consistency"],"backgroundTag":"tiling-consistency-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}