Colab Runtime Error - Jaxlib version mismatch (strikes again)

[Robbe Wulgaert]

Hi all

When running the Colab Notebook, at the import step, the Jax and Jaxlib modules are causing an error.

RuntimeError: jaxlib is version 0.3.25, but this version of jax requires version >= 0.4.2.

We've had a similar issue around the 12th of September with Jax, Jaxlib and Flax. Probably an update in Colab or one of those dependencies has made this resurface.

In the meantime, users who are experiencing issues can circumvent this error using following janky solution:

!pip install -qqq jax==0.4.2
!pip install -qqq jaxlib==0.4.2

Limitation: This will run the beam search on the CPU of the Colab, thus increasing inference times.
You can make the GPU kick in again by using instead:

!pip install -qqq jax==0.4.2
!pip install -qqq jaxlib==0.4.2
!pip install -qqq "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

I have used the following examples to test inference using the aforementioned janky solutions:

ωρανια.   αστ?ρας ηρε?νησα σο??ι φρ?νι π??ρι τε εοι?ος   ουν?μα εχ?. λε??μαι δε η δι?ς ουρα?ιη.   ονε?του.

and

λ??ουσιν ἃ θέλου??ν λεγέ??σαν οὐ μέλι ??ι σὺ φίλι με συνφέ?ι σοι

Kind regards!

[Yannis Assael]

Robbe, our team would like to share a big thank you for all your help and support! We've updated the minimum JAX version and the problem should be now resolved.