Colab TPU runtime requires jax<=0.3.25

[andsteing]

The following fails on Colab TPU runtime:

>>> !pip install flax
>>> import jax
RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.7.

[andsteing]

As mentioned in issues #2950 and #2861, the Google Colab TPU runtime does not support new JAX versions anymore.

There are some technical roadblocks in updating the Colab TPU runtime infra, and we might be blocked for the foreseeable future.

Until then, it's only possible to use these runtimes with an old pinned version of JAX:

!pip install jax==0.3.25 jaxlib==0.3.25 flax

[andsteing]

BTW this is a useful Colab for testing the setup on TPU runtimes:

https://colab.research.google.com/github/google/flax/blob/main/tests/colab_tpu_jax_version.ipynb