[shard-map] in eager shmap, handle all rep rule output cases #27797
Conversation
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
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
By convention, rep_rules can return three kinds of thing: 1. a sequence (tuple or list), 2. a single set, or 3. a single None. Even rules for primitives with multiple results can return single objects rather than sequences; the reason is that it's convenient not ot have to infer the number of outputs for higher-order primitives. In the latter two cases we rely on the caller (in this case, ShardMapTrace.process_primitive) to 'broadcast' the singleton result to a list of results equal to the number of outputs. Previously, the code was checking `if type(out_rep) is set`, which doesn't handle case 3. (We briefly tried another fix direction where we don't allow case 3, because we don't have case 3 in the upcoming VMA type system which replaces this stuff. But until that lands the easiest fix is just to handle all cases correctly.) fixes jax-ml#26148, fixes jax-ml#27673 Co-authored-by: Justin Fu <justinfu@google.com>
justinjfu
approved these changes
Apr 7, 2025
google-ml-butler
bot
added
kokoro:force-run
pull ready
mattjj
added a commit
to mattjj/jax
that referenced
this pull request
Apr 9, 2025
…k_rep this was essentially another instance of the jax-ml#27797 fix fixes jax-ml#24762
charleshofer
pushed a commit
to ROCm/jax
that referenced
this pull request
Apr 30, 2025
…k_rep this was essentially another instance of the jax-ml#27797 fix fixes jax-ml#24762
charleshofer
pushed a commit
to ROCm/jax
that referenced
this pull request
May 1, 2025
…k_rep this was essentially another instance of the jax-ml#27797 fix fixes jax-ml#24762
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
By convention, rep_rules can return three kinds of thing:
Even rules for primitives with multiple results can return single objects rather than sequences; the reason is that it's convenient not ot have to infer the number of outputs for higher-order primitives.
In the latter two cases we rely on the caller (in this case, ShardMapTrace.process_primitive) to 'broadcast' the singleton result to a list of results equal to the number of outputs.
Previously, the code was checking
if type(out_rep) is set, which doesn't handle case 3.(We briefly tried another fix direction where we don't allow case 3, because we don't have case 3 in the upcoming VMA type system which replaces this stuff. But until that lands the easiest fix is just to handle all cases correctly.)
fixes #26148, fixes #27673
This PR replaces #27720