TPU Installation broken because of change in Orbax

[akhilkedia]

Hi!

Running the scripts/tpu_vm_setup.sh file currently errors out because of a change in the orbax package, as discussed here google/orbax#436

Collecting orbax
  Using cached orbax-0.1.8.tar.gz (1.6 kB)
  Preparing metadata (setup.py) ... error
  error: subprocess-exited-with-error
  
  × python setup.py egg_info did not run successfully.
  │ exit code: 1
  ╰─> [3 lines of output]
      
      *** Orbax is a namespace, and not a standalone package. For model checkpointing and exporting utilities, please install `orbax-checkpoint` and `orbax-export` respectively (instead of `orbax`). ***
      
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
error: metadata-generation-failed

× Encountered error while generating package metadata.
╰─> See above for output.

The orbax dependency originates from flax, and flax==0.6.8 requires orbax

While latest flax has changed this to orbax-checkpoint (as can be seen by running curl -L 'https://pypi.python.org/pypi/flax/json' | jq '.info.requires_dist'), the version of flax in scripts/tpu_vm_setup.sh directly depends on orbax.

So perhaps the flax dependency version should be changed in scripts/tpu_vm_setup.sh, or orbax==0.1.7 should be added to fix the dependency.

[young-geng]

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.