{"record":{"id":"f8a4bf6f80915e49","repo":"jax-ml/jax","slug":"invalid-in-out-transforms-in-transform-in-trans","errorCode":null,"errorMessage":"Invalid in/out transforms. In transform: {in_transform}, out transform: {transform}","messagePattern":"Invalid in/out transforms\\. In transform: (.+?), out transform: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2180,"sourceCode":"        in_transform, lc.TileTransform\n    ):\n      raise NotImplementedError(\n          f\"Unsupported in/out transforms. In transform: {in_transform}, out\"\n          f\" transform: {transform}\"\n      )\n    tiling_len = len(in_transform.tiling)\n    tiling_offset = len(permutation) - tiling_len\n    if any(dim < tiling_offset for dim in permutation[-tiling_len :]):\n      raise ValueError(\n          f\"Cannot tile a transpose ({permutation}). Tiling dims\"\n          f\" ({permutation[-tiling_len:]}) cannot contain non-tiled dims.\"\n          f\" All of them must be >= {tiling_offset}.\"\n      )\n    dims = [-1] * len(permutation)\n    dims = dims[:-tiling_len] + list(in_transform.tiling)\n    permuted_dims = tuple(dims[permutation[i]] for i in range(len(dims)))\n    if permuted_dims[-tiling_len:] != transform.tiling:\n      raise ValueError(\n          f\"Invalid in/out transforms. In transform: {in_transform}, out\"\n          f\" transform: {transform}\"\n      )\n    new_permutation = permutation + [\n        x + tiling_len for x in permutation[-tiling_len:]\n    ]\n    new_permutation = _permutation_to_affine_map_attr(new_permutation)\n\n  new_transpose_op = memref.TransposeOp(\n      transform_type(ir.MemRefType(op.result.type), out_transforms),\n      unwrapped_in_ref,\n      new_permutation,\n  )\n\n  out_transforms = inference_utils.out_transforms(op)[0]\n  wrapped_ref = wrap_transformed_memref(\n      new_transpose_op.result, op.result.type, out_transforms\n  )","sourceCodeStart":2162,"sourceCodeEnd":2198,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2162-L2198","documentation":"After applying the permutation to the input transform's tiling, the result must equal the output transform's declared tiling. This error means the annotated output tiling is not the permutation of the input tiling as the transpose requires.","triggerScenarios":"memref.transpose with transforms where permuted input tiling != output transform's tiling, e.g. in tiling (2,4) and permutation that maps it to (4,2) but out tiling declared (2,4).","commonSituations":"Hand-specifying the output transform without deriving it from the permutation; inference usually computes it for you when you omit the annotation.","solutions":["Let transform inference compute the output tiling instead of annotating it manually","Or set out tiling = tuple(dims[p[i]] ... ) i.e. the input tiling permuted by the permutation","Verify by applying the permutation to in_transform.tiling and comparing to out_transform.tiling before emitting the op"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"dims = [-1]*len(perm)\ndims = dims[:-tl] + list(in_t.tiling)\nexpected = tuple(dims[perm[i]] for i in range(len(dims)))[-tl:]\nassert expected == out_t.tiling, f'expected out tiling {expected}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Let inference compute output tilings","Derive out tiling by permuting the in tiling"],"tags":["jax","mosaic-gpu","transpose","tiling","consistency"],"backgroundTag":"tiling-consistency-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}