Lower shard_map in partial explicit manual mode correctly by using the aval.sharding.spec instead of `in_specs` since the former contain more info (explicit mesh axes too) while in_specs only refer to axes passed in axis_names. #40776

Merged
copybara-service[bot] opened 3:38pm on September 17, 2026 wanted to merge 151 Δ into google/jax main from
test_983407126

Pull request #40776 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 151
Maximum comment count per PR 10 0
Pull request maximum business days between activity 3 1
Pull request maximum business days open 5 1

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/shard_map.py 7 12 2266
tests/shard_map_test.py 0 44 6236