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

[IronCraftMan]

Whenever I try to run via docker "docker-compose up" in the main directory, I get this output on loop.

--> Starting DALL-E Server. This might take up to two minutes. dalle-backend | Traceback (most recent call last): dalle-backend | File "app.py", line 18, in dalle-backend | from dalle_model import DalleModel dalle-backend | File "/app/dalle_model.py", line 5, in dalle-backend | import jax dalle-backend | File "/usr/local/lib/python3.8/dist-packages/jax/init.py", line 116, in dalle-backend | from .experimental.maps import soft_pmap as soft_pmap dalle-backend | File "/usr/local/lib/python3.8/dist-packages/jax/experimental/maps.py", line 26, in dalle-backend | from .. import numpy as jnp dalle-backend | File "/usr/local/lib/python3.8/dist-packages/jax/numpy/init.py", line 19, in dalle-backend | from . import fft as fft dalle-backend | File "/usr/local/lib/python3.8/dist-packages/jax/numpy/fft.py", line 17, in dalle-backend | from jax._src.numpy.fft import ( dalle-backend | File "/usr/local/lib/python3.8/dist-packages/jax/_src/numpy/fft.py", line 19, in dalle-backend | from jax import lax dalle-backend | File "/usr/local/lib/python3.8/dist-packages/jax/lax/init.py", line 332, in dalle-backend | from jax._src.lax.fft import ( dalle-backend | File "/usr/local/lib/python3.8/dist-packages/jax/_src/lax/fft.py", line 145, in dalle-backend | xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft dalle-backend | AttributeError: module 'jaxlib.pocketfft' has no attribute 'pocketfft' dalle-backend exited with code 1

[centipeda]

@IronCraftMan The fix in #79 by @Marcelovk worked for me to fix this.

[cdmichaelb]

@IronCraftMan The fix in #79 by @Marcelovk worked for me to fix this.

This fix did not work for me, can you tell me what to do in case I am doing it wrong?

[centipeda]

I changed the one line in backend/Dockerfile to the one in the merge request. I think the exact version of jax to install might depend on your graphics card situation, so I'm not sure if it'll work for everyone.

[cdmichaelb]

Ah ok, I tried that too. I have an RTX 3080, do you?

I changed the one line in backend/Dockerfile to the one in the merge request. I think the exact version of jax to install might depend on your graphics card situation, so I'm not sure if it'll work for everyone.

[Marcelovk]

Just checking - are you sure that you are doing a docker-compose build again to pick the changes?

[centipeda]

I have a 2070 Super. And yes, don't forget to docker-compose build after making the changes.

[cdmichaelb]

@centipeda @Marcelovk Thanks, docker-compose build was what I was missing. All seems to be good now (or at least it's booting!)