cannot import name 'ShapedArray' from 'jax'

[VictorHavard]

After installing all of the prerequisites including installing jax and testing it, when I create the venv and attempt to run the server I get "cannot import name 'ShapedArray' from 'jax'".

ShapedArray was removed in jax in version 0.4.14. Current latest version of jax is 0.4.23.

Do I need to install the older version of jax or is there some way around this?

[VictorHavard]

Works when using the docker based version because it specifies to use 0.4.13.

[koiosjojo]

Changing the requirements file found under dalle-playground to the following and reinitializing with pip fixed the issue for me:
Flask==2.1.0
Flask-Cors==3.0.9
Flask-RESTful==0.3.8
Werkzeug==2.2.2
jax==0.4.13
jaxlib==0.4.13
flax
rich
tqdm
flask_cloudflared
accelerate
transformers
diffusers