Streamlit app fails to load

[Kaushalya Madhawa]

Stacktrace:

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

File "/home/user/.local/lib/python3.8/site-packages/streamlit/script_runner.py", line 354, in _run_script
    exec(code, module.__dict__)
File "/home/user/app/app.py", line 7, in <module>
    from medclip.modeling_hybrid_clip import FlaxHybridCLIP
File "/home/user/app/medclip/modeling_hybrid_clip.py", line 18, in <module>
    import flax.linen as nn
File "/home/user/.local/lib/python3.8/site-packages/flax/__init__.py", line 36, in <module>
    from . import core
File "/home/user/.local/lib/python3.8/site-packages/flax/core/__init__.py", line 15, in <module>
    from .axes_scan import broadcast
File "/home/user/.local/lib/python3.8/site-packages/flax/core/axes_scan.py", line 17, in <module>
    import jax
File "/home/user/.local/lib/python3.8/site-packages/jax/__init__.py", line 108, in <module>
    from .experimental.maps import soft_pmap
File "/home/user/.local/lib/python3.8/site-packages/jax/experimental/maps.py", line 25, in <module>
    from .. import numpy as jnp
File "/home/user/.local/lib/python3.8/site-packages/jax/numpy/__init__.py", line 16, in <module>
    from . import fft
File "/home/user/.local/lib/python3.8/site-packages/jax/numpy/fft.py", line 17, in <module>
    from jax._src.numpy.fft import (
File "/home/user/.local/lib/python3.8/site-packages/jax/_src/numpy/fft.py", line 19, in <module>
    from jax import lax
File "/home/user/.local/lib/python3.8/site-packages/jax/lax/__init__.py", line 331, in <module>
    from jax._src.lax.fft import (
File "/home/user/.local/lib/python3.8/site-packages/jax/_src/lax/fft.py", line 144, in <module>
    xla.backend_specific_translations['cpu'][fft_p] = pocketfft.pocketfft