Colab Error - AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft'

[Robbe Wulgaert]

Hi,

I noticed an error propping up in the Colab referenced in this GitHub page, and in the version I've been using with my student.
The error originates in the "Imports" code cell of the reference Colab Notebook and is (probably) due to an upgrade in the Colab dependencies.

The error occurs when running:

import flax

The error:

AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft' 

Solution (a bit janky but seems to work):

!pip install jax==0.2.21
!pip install jaxlib==0.1.69

I still have to do some speed tests to check if this approach has no adverse effects.

[Robbe Wulgaert]

Adding to the above 'janky solution':
Input:

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

Time required for inference: 10 min 09 sec. This is due to the system not utilising CUDA. Less the ideal performance.

Proposed adjustment when encountering following error in a recent Google Colab-instance:

AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft' 

Alternatieve, less janky, solution (uses CUDA, dropping inference time back to normal levels, 3 min 29 sec):

!pip install --upgrade jax==0.2.21 jaxlib==0.1.69+cuda111 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[Yannis Assael]

Hi Robbe, thank you so much for bringing this to our attention, providing a solution, and for all your fantastic work using Ithaca. We are looking into the issue and will soon provide a solution.

[Yannis Assael]

Hi Robbe, I think it should be fixed now, thanks again!