Installation Issue from flax vs jax compatibility

[david2611]

Following the install instructions on the main README page caused an error seemingly due to incompatibilities between jax and flax.

Following install instructions gave jax version 0.4.16 and flax version 0.5.3

Running the example training code gave the following error ImportError: cannot import name 'ShapedArray' from 'jax'

I have not used jax or flax before but from doing some digging this seems to be because jax hasn't provided ShapedArray since version 0.4.14.

Upgrading flax to version 0.7.0 got past this problem but opened up a new issue of flax.optim no longer being a part of the flax library.

What is the recommended solution to this problem and can code/documentation be updated to reflect this compatibility issue?

[fasogbon]

Same issue here

[fasogbon]

Pleas help to solve this issue! Thank you

[fasogbon]

This worked for me using python 3.6:
pip3 install --upgrade jax==0.2.14 jaxlib==0.1.69+cuda111 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[KaneLindsay]

I got a working Conda environment with Python 3.9, Ubuntu 20.04, CUDA Toolkit 11.8 and cuDNN 8.6 by specifying the requirement versions in jax3d/requirements.txt as:

chex
clu
einops
flax
gin-config
imageio
jax[cuda11_local] == 0.4.4
jaxlib @ https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.4.4+cuda11.cudnn86-cp39-cp39-manylinux2014_x86_64.whl
mediapy
numpy
pandas
simple_parsing
tensorflow[and-cuda]
tfds-nightly
typing_extensions
xmanager
scikit-image 

sunds @ git+https://github.com/google-research/sunds
etils[array_types,enp,epath,epy,etqdm,etree] @ git+https://github.com/google/etils

Then,

pip install flax==0.5.3
pip install --upgrade "jax3d[nesf]"