[Pallas:MGPU] Support transposing the swizzled dimension in `UnswizzleRef`. #41854

Closed
copybara-service[bot] opened 11:28am on October 21, 2026 wanted to merge 226 Δ into google/jax main from
test_1003935315

Pull request #41854 quality checks

Code files and functions touched by this pull request are analyzed below.

Quality targets

How does this pull request compare to chosen code quality targets? Read more about code quality targets.

Target name Target value This PR
Maximum Diff Delta per PR 500 226
Maximum comment count per PR 10 0
Pull request maximum business days between activity 3 16
Pull request maximum business days open 5 17

Coupled files

Are there any files from this commit that are always committed alongside another file? This often suggests that there exists code coupling to potentially remove. Read more about code coupling

No file coupling detected. 🎉

Biggest files changes in pull request

Path View whole file Lines removed Lines added Final line count
jax/_src/pallas/mosaic_gpu/core.py 16 60 2267
tests/pallas/mosaic_gpu_test.py 0 182 11124

️ Biggest functions grown or shrunk by pull request

Which functions had their line count change as part of this pull request?

Qualified function name View whole function Lines removed Lines added Final line count
UnswizzleRef::commute_ndindexer 0 19 38
UnswizzleRefTest::test_commute_ndindexer 0 25 13
UnswizzleRefTest::test_commute_transpose 0 15 13
UnswizzleRefTest::test_commute_ndindexer::test_commute_ndindexer_sliced_swizzled_dim_raises 0 9 13
UnswizzleRef::commute_reshape 5 4 13
UnswizzleRefTest::test_commute_ndindexer_int_index_swizzled_dim_raises 0 13 10
UnswizzleRef::__post_init__ 2 10 10
UnswizzleRefTest::test_commute_transpose::test_commute_reshape_transposed_swizzled_dim_raises 0 11 9
UnswizzleRefTest::test_commute_transpose::test_commute_reshape_wrong_minor_dim_raises 0 10 8
UnswizzleRefTest::test_commute_transpose::test_commute_reshape 0 10 8
UnswizzleRefTest::test_undo_transposed_swizzled_dim_raises 0 10 8
UnswizzleRefTest::test_transform_type_invalid_shape_raises 0 10 8
UnswizzleRef::commute_transpose 5 3 7
UnswizzleRefTest::test_transform_type_valid 0 8 6
UnswizzleRef::swizzle_elems 0 3 6
UnswizzleRefTest::test_commute_ndindexer::test_commute_ndindexer_sliced_swizzled_dim_raises::f 0 6 4
UnswizzleRefTest::test_undo 0 6 4
UnswizzleRefTest::test_non_negative_swizzled_dim_raises 0 6 4
UnswizzleRefTest::test_invalid_swizzle_raises 0 6 4