Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixes incorrect sharding lowering for key<fry> in shard_map partial-auto #25692

Merged
merged 1 commit into from
Dec 30, 2024

Conversation

yliu120
Copy link
Contributor

@yliu120 yliu120 commented Dec 29, 2024

The extended type key<fry> needs to be expanded to u32[2] before being used in manual_proto, otherwise the ranks in the sharding is not correct.

@yliu120
Copy link
Contributor Author

yliu120 commented Dec 29, 2024

@mattjj Here is another bug. Please take a look. Thanks.

Copy link
Collaborator

@mattjj mattjj left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Dec 30, 2024
@copybara-service copybara-service bot merged commit 494c157 into jax-ml:main Dec 30, 2024
25 checks passed
@yliu120 yliu120 deleted the rng_key_sharding branch December 30, 2024 18:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants