maxtext
maxtext copied to clipboard
Add a new script to build the Jax stable stack based image for xpk and helm
Description
The Jax stable stack image contain Jax, cuda and MaxText, but lack of some gcp deps. This PR intends to add these deps and build an image that works for xpk and helm.
Tests
Manually tested, the workflow and image works.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
- [x] I have performed a self-review of my code.
- [x] I have necessary comments in my code, particularly in hard-to-understand areas.
- [x] I have run end-to-end tests tests and provided workload links above if applicable.
- [x] I have made or will make corresponding changes to the doc if needed.
This PR has been automatically marked as stale because it has not had recent activity. It will be closed soon if no further activity occurs. Thank you for your contributions.
This PR was closed because it has been inactive for a while. Please reopen it if you are still working on it.