Compatibility issues about flax/jax/optim/optax

[AnabelicSun]

While running the example code in Readme, such error occurs:

./run.sh
2024-10-24 14:45:56.243665: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2024-10-24 14:45:56.301708: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2024-10-24 14:45:56.318591: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
Traceback (most recent call last):
File "/home/anabelian/GIT/meliad/transformer/ht_main.py", line 28, in
from transformer import launcher
File "/home/anabelian/GIT/meliad/transformer/launcher.py", line 23, in
import training_loop
File "/home/anabelian/GIT/meliad/training_loop.py", line 32, in
import optimizer_config as opt_config
File "/home/anabelian/GIT/meliad/optimizer_config.py", line 21, in
from flax import optim
ImportError: cannot import name 'optim' from 'flax' (/home/anabelian/alphageo/lib/python3.11/site-packages/flax/init.py)

it seems that newer version of flax does not support optim, however in requirement:
flax>=0.5.0
jax>=0.3.13
optax>=0.1.2

it seems that it is better to specify flax = 0.5.3 and jax = 0.3.13 to prevent problems?