Open david2611 opened 11 months ago
Same issue here
Pleas help to solve this issue! Thank you
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
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]"
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?