{"record":{"id":"89c16b84f1d63f49","repo":"jax-ml/jax","slug":"swizzle-mismatch-in-transforms-swizzle-in-swizz","errorCode":null,"errorMessage":"Swizzle mismatch. In transforms swizzle: {in_swizzle}, out transforms swizzle {out_swizzle}.","messagePattern":"Swizzle mismatch\\. In transforms swizzle: (.+?), out transforms swizzle (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2135,"sourceCode":"  return ir.AffineMapAttr.get(ir.AffineMap.get_permutation(permutation))\n\n\n@_register_lowering(memref.TransposeOp, support_warp_semantics=True)\ndef _memref_transpose_op_lowering_rule(\n    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:","sourceCodeStart":2117,"sourceCodeEnd":2153,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2117-L2153","documentation":"The transpose lowering requires the swizzle component of the input and output transform annotations to be identical. Transpose permutes dimensions but cannot re-derive a different swizzle pattern, so mismatched swizzles raise this ValueError.","triggerScenarios":"memref.transpose with in_transforms and out_transforms whose swizzle_from_transforms_attr values differ (e.g. input swizzle=SWIZZLE_32B, output swizzle=SWIZZLE_128B or None).","commonSituations":"Transposing a swizzled shared-memory tile while also trying to change its swizzle mode in one op; common when tuning shared-memory layouts for MMA pipelines.","solutions":["Keep the same swizzle on both sides of the transpose; apply a separate op to change swizzle","Drop swizzle from both annotations if the layout doesn't need it","Check swizzle_from_transforms_attr on both attrs and align them before building the transpose"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert swizzle_from_transforms_attr(in_attr) == swizzle_from_transforms_attr(out_attr)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep swizzle identical across transpose operands","Change swizzle with a dedicated op"],"tags":["jax","mosaic-gpu","transpose","swizzle","transforms"],"backgroundTag":"swizzle-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}