Home » Oreps » Jax » All Pull Requests » #40776
google/jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
Watch
1
Commit Activity
Pull Requests
Directory Browser
Velocity History
Releases

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

Compare code on GitHub
Merged
copybara-service[bot] opened 3:38pm on September 17, 2026 wanted to merge 151 Ξ” into google/jax main from
test_983407126
copybara-service[bot]
google/jax
test_983407126
Diff Delta:
151
About 70 Diff Delta/hour
Classified as:  Feature Add

copybara-service[bot]'s Description of Work

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.

Also add partial manual support to the check which checks if in_specs == aval.sharding.spec (otherwise it errors out).

2 total changed files
Loading changes...
End of PR:
Like the credits, but with fewer stunts.
GitClear uses cookies to ensure you get the best experience on our website. Learn more about our cookie policy