{"record":{"id":"08b2d4fbb7e09746","repo":"jax-ml/jax","slug":"size-mismatch-for-in-out-transforms-in-transforms","errorCode":null,"errorMessage":"Size mismatch for in/out transforms. In transforms: {in_transforms}, out transforms: {out_transforms}.","messagePattern":"Size mismatch for in/out transforms\\. In transforms: (.+?), out transforms: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2140,"sourceCode":"    ctx: LoweringContext, op: memref.TransposeOp\n) -> Sequence[ir.Value]:\n  del ctx\n\n  in_transforms_attr = inference_utils.in_transforms(op)[0]\n  unwrapped_in_ref = unwrap_transformed_memref(op.in_, in_transforms_attr)\n  in_swizzle = swizzle_from_transforms_attr(in_transforms_attr)\n  in_transforms = memref_transforms_from_transforms_attr(in_transforms_attr)\n  out_transforms_attr = inference_utils.out_transforms(op)[0]\n  out_swizzle = swizzle_from_transforms_attr(out_transforms_attr)\n  out_transforms = memref_transforms_from_transforms_attr(out_transforms_attr)\n\n  if in_swizzle != out_swizzle:\n    raise ValueError(\n        f\"Swizzle mismatch. In transforms swizzle: {in_swizzle}, out transforms\"\n        f\" swizzle {out_swizzle}.\"\n    )\n  if len(out_transforms) != len(in_transforms):\n    raise ValueError(\n        f\"Size mismatch for in/out transforms. In transforms: {in_transforms},\"\n        f\" out transforms: {out_transforms}.\"\n    )\n  if not out_transforms:\n    new_permutation = op.permutation\n  else:\n    permutation = [\n        ir.AffineDimExpr(e).position\n        for e in op.permutation.value.results\n    ]\n    # We expect to have the same transforms on in/out, up to permutation of the\n    # out transforms.\n    # For example, for 3D input:\n    #   permutation: (0, 2, 1)\n    #   in_transforms: TilingTransform((32, 8))\n    # We expect:\n    #   out_transforms: TilingTransform((8, 32))\n    # TODO(olechwierowicz): Support multiple transforms.","sourceCodeStart":2122,"sourceCodeEnd":2158,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2122-L2158","documentation":"memref.transpose lowering requires the number of MemRefTransforms on input and output to be equal. Extra or missing transforms on either side cannot be reconciled by permutation alone.","triggerScenarios":"Building a transpose where in_transforms has, say, one tile transform but out_transforms has zero or two (or vice versa).","commonSituations":"Transform inference producing a default empty transform on one side while the other is explicitly annotated; typically from manually constructing the op rather than using Mosaic's Python builder APIs.","solutions":["Set identical transform lists (same count) on input and output of the transpose","Use Mosaic's high-level transpose/permute helper so inference fills both sides consistently","If no transforms are intended, clear the annotations on both sides"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert len(in_transforms) == len(out_transforms), 'transform count mismatch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Rely on inference rather than manual annotations","Mirror transform counts on both sides"],"tags":["jax","mosaic-gpu","transpose","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"}