google/jax
copybara-service[bot]'s Description of Work
[Pallas:MGPU] Support transposing the swizzled dimension in UnswizzleRef.
Previously, UnswizzleRef assumed the swizzled dimension was always the minormost dimension (-1) and rejected any TransposeTransform that moved it. Because of this restriction, _get_lowering_rule and _swap_lowering_rule had to manually pop UnswizzleRef before calling _handle_transforms: when a tiled and swizzled SMEM reference is transposed, commuting TransposeTransform past UntilingTransform also transposes the 2D tile dimensions, which moves the swizzled dimension away from -1.
To allow UnswizzleRef to flow through _handle_transforms uniformly with UntilingTransform (in a followup change), UnswizzleRef needs to track which dimension is swizzled (swizzled_dim) across TransposeTransform and NDIndexer rather than assuming it is always -1, while enforcing that operations requiring the swizzled dimension to remain minormost (undo and commute_reshape) or un-sliced (commute_ndindexer) check swizzled_dim.