{"record":{"id":"bb118ce2757f567c","repo":"jax-ml/jax","slug":"unsupported-in-out-transforms-in-transform-in-t","errorCode":null,"errorMessage":"Unsupported in/out transforms. In transform: {in_transform}, out transform: {transform}","messagePattern":"Unsupported in/out transforms\\. In transform: (.+?), out transform: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2164,"sourceCode":"  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.\n    [transform] = out_transforms\n    [in_transform] = in_transforms\n    if not isinstance(transform, lc.TileTransform) or not isinstance(\n        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}\"","sourceCodeStart":2146,"sourceCodeEnd":2182,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2146-L2182","documentation":"When a transpose carries transforms, both the input and output transform must be TileTransforms (only tiling is supported, not swizzle or other transform kinds, and only a single one each).","triggerScenarios":"memref.transpose with non-empty transforms where either in_transform or out_transform is not an instance of lc.TileTransform (e.g. a swizzle-only or slice transform).","commonSituations":"Stacking a swizzle annotation onto a transposed tiled tensor; hitting the documented TODO(olechwierowicz) limitation of single-tile-transform support.","solutions":["Reduce transforms on both sides to a single TileTransform each","Move swizzle/other transforms outside the transpose","If you need multiple transforms around a transpose, materialize them as explicit reshape/layout ops instead"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def both_single_tile_transforms(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":["Use only TileTransforms around transposes","Materialize swizzles outside the transpose"],"tags":["jax","mosaic-gpu","transpose","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"}