-
Notifications
You must be signed in to change notification settings - Fork 233
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
TPU Installation broken because of change in Orbax #84
Comments
We've also noticed that recently, and we think that this would be a good time for us to upgrade the JAX and Flax version. We are still testing some code to make sure everything is compatible with the newest version of JAX and Flax, and will release it as soon as possible. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Hi!
Running the
scripts/tpu_vm_setup.sh
file currently errors out because of a change in theorbax
package, as discussed here google/orbax#436The orbax dependency originates from
flax
, andflax==0.6.8
requiresorbax
While latest flax has changed this to
orbax-checkpoint
(as can be seen by runningcurl -L 'https://pypi.python.org/pypi/flax/json' | jq '.info.requires_dist'
), the version of flax inscripts/tpu_vm_setup.sh
directly depends on orbax.So perhaps the flax dependency version should be changed in
scripts/tpu_vm_setup.sh
, ororbax==0.1.7
should be added to fix the dependency.The text was updated successfully, but these errors were encountered: