{"record":{"id":"803721ae79f5ef92","repo":"jax-ml/jax","slug":"expected-the-same-number-of-in-out-transforms-but","errorCode":null,"errorMessage":"Expected the same number of in/out transforms, but got {in_transforms=} and {out_transforms=}","messagePattern":"Expected the same number of in/out transforms, but got (.+?) and (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2254,"sourceCode":"      output_shape=op.output_shape,\n      static_output_shape=out_transformed_ty.shape,\n  )\n\n  wrapped_ref = wrap_transformed_memref(\n      new_expand_shape_op.result, op.result.type, out_transforms\n  )\n  return [wrapped_ref]\n\n\n# TODO(bchetioui): find a way to consolidate the logic that is shared logic with\n# 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(","sourceCodeStart":2236,"sourceCodeEnd":2272,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2236-L2272","documentation":"memref.collapse_shape lowering requires exactly the same number of input and output transforms. A mismatch means the layout annotations are inconsistent and the collapse semantics are undefined.","triggerScenarios":"memref.collapse_shape where len(in_transforms) != len(out_transforms), e.g. input annotated with one tiling and output with none.","commonSituations":"Collapsing a tiled smem buffer while manually constructing the result type without carrying the transform annotation over.","solutions":["Mirror the transform count: same number of transforms (typically one TileTransform) on both sides","Omit manual annotations and rely on inference for the output","For untransformed memrefs ensure both annotations are empty"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert len(in_transforms) == len(out_transforms)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Mirror transform counts on collapse operands","Prefer inferred output transforms"],"tags":["jax","mosaic-gpu","collapse-shape","transforms","count-mismatch"],"backgroundTag":"unsupported-layout-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}