Generalize out_sharding to work with NamedSharding and Format#5246
Merged
copybara-service[bot] merged 2 commits intogoogle:mainfrom Mar 20, 2026
Merged
Conversation
Contributor
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
samanklesaria
commented
Feb 13, 2026
|
|
||
| def get_pspec(sharding, sharding_rules = None) -> PartitionSpec: | ||
| """Given an `nnx.Variable`, return its `PartitionSpec`.""" | ||
| return apply_rules(sharding, sharding_rules) # type: ignore |
Collaborator
Author
There was a problem hiding this comment.
I'm maintaining this function for backwards compat, but because it now works on NamedShardings and Formats too, we can't have the name be called 'get_pspec'. What this function really does it apply a mapping transformation to a sharding. So I've called the more general version 'apply_rules'.
Collaborator
There was a problem hiding this comment.
yeah maybe we can call it get_sharding or something? Seems like the return type also doesn't match.
cgarciae
reviewed
Mar 4, 2026
cgarciae
reviewed
Mar 4, 2026
samanklesaria
force-pushed
the
generalize_out_sharding
branch
from
52ce675 to
1c2cec4
Compare
cgarciae
approved these changes
Mar 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Like the reverted PR #5089, this PR generalized flax's spmd APIs to allow sharding_metadata to be NamedSharding or Format objects.