Conversation
qihqi
approved these changes
Jul 7, 2025
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.
Given the decision to sync with a stable release, modifies the update_deps.py script to optionally determine the latest release of Jax and get the corresponding openxla hash and libtpu version.
Running the script changes the XLA hash from 9084478 to 3d5ece6. The latter is the XLA commit pinned by Jax 0.6.2, which is what we want. The former was the last commit on the day of the release. In this case they both work, but it's a nice verification that this script can get us a more compatible version with less effort.
Note that setup.py has a
USE_NIGHTLYbool which, if False, will pull jax and jaxlib using their version numbers instead of pulling the corresponding nightly from the package registries. This seems like a change we want to make but doing so in this case fails withThis failure is because we are missing symbol
sym_is_non_overlapping_and_dense_custom, introduced a few days ago here. So we're using an older, stable version of pytorch. I don't understand why this is the case, based on setup.py. But in the meantime I keepUSE_NIGHTLYas True. This means we're still technically pulling a nightly version, but it's the nightly built from the most recent stable Jax version.