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

Replace pjit with jax.jit by converting PartitionSpec to NamedShardings. #193

Merged
merged 1 commit into from
Oct 4, 2023

Conversation

michelle-yooh
Copy link
Collaborator

…onSpec to NamedShardings.

@michelle-yooh michelle-yooh requested a review from rwitten as a code owner October 2, 2023 20:56
@michelle-yooh michelle-yooh changed the title Use jax.jit in Maxtext and remove pjit. Do this by converting Partiti… Replace pjit with jax.jit by converting PartitionSpec to NamedShardings. Oct 2, 2023
@michelle-yooh michelle-yooh force-pushed the yooh-jit-with-testing branch from 7ce07ab to fbc96c2 Compare October 4, 2023 21:15
… NamedShardings.

Update unit test for decode.py
@michelle-yooh michelle-yooh force-pushed the yooh-jit-with-testing branch from 1a80078 to 688c698 Compare October 4, 2023 22:54
Copy link
Collaborator

@khatwanimohit khatwanimohit left a comment

Choose a reason for hiding this comment

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

LGTM!

@michelle-yooh michelle-yooh merged commit f9c8e6d into main Oct 4, 2023
A9isha pushed a commit that referenced this pull request Apr 11, 2024
* fix a flaky wget issue.

* fix a dir issue.

* fix comments

* fix missing dir
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants