New python venv fails to pip install the flax mnist example requirements about flax HOT 3 OPEN

python -m venv env_imagenet
source env_imagenet/bin/activate
unset LD_LIBRARY_PATH
pip install -U "jax[cuda12]"
cd flax
pip install -e .
cd examples/imagenet/
pip install -r requirements.txt

Package                      Version
---------------------------- -----------
absl-py                      1.0.0
astunparse                   1.6.3
attrs                        23.2.0
cachetools                   5.3.3
certifi                      2024.6.2
charset-normalizer           3.3.2
chex                         0.1.7
clu                          0.0.6
contextlib2                  21.6.0
contourpy                    1.2.1
cycler                       0.12.1
dill                         0.3.8
dm-tree                      0.1.8
etils                        1.7.0
flatbuffers                  24.3.25
flax                         0.6.5
fonttools                    4.53.0
fsspec                       2024.6.0
future                       1.0.0
gast                         0.4.0
google-auth                  2.30.0
google-auth-oauthlib         0.4.6
google-pasta                 0.2.0
googleapis-common-protos     1.63.1
grpcio                       1.64.1
h5py                         3.11.0
idna                         3.7
importlib_resources          6.4.0
jax                          0.4.29
jax-cuda12-pjrt              0.4.29
jax-cuda12-plugin            0.4.29
jaxlib                       0.4.29
keras                        2.11.0
kiwisolver                   1.4.5
libclang                     18.1.1
Markdown                     3.6
markdown-it-py               3.0.0
MarkupSafe                   2.1.5
matplotlib                   3.8.4
mdurl                        0.1.2
ml-collections               0.1.0
ml-dtypes                    0.4.0
msgpack                      1.0.8
nest-asyncio                 1.6.0
numpy                        1.22.0
nvidia-cublas-cu12           12.5.2.13
nvidia-cuda-cupti-cu12       12.5.39
nvidia-cuda-nvcc-cu12        12.5.40
nvidia-cuda-runtime-cu12     12.5.39
nvidia-cudnn-cu12            9.1.1.17
nvidia-cufft-cu12            11.2.3.18
nvidia-cusolver-cu12         11.6.2.40
nvidia-cusparse-cu12         12.4.1.24
nvidia-nccl-cu12             2.21.5
nvidia-nvjitlink-cu12        12.5.40
oauthlib                     3.2.2
opt-einsum                   3.3.0
optax                        0.1.3
orbax                        0.1.9
orbax-checkpoint             0.5.16
packaging                    24.1
pillow                       10.3.0
pip                          22.0.2
promise                      2.3
protobuf                     3.19.6
pyasn1                       0.6.0
pyasn1_modules               0.4.0
Pygments                     2.18.0
pyparsing                    3.1.2
python-dateutil              2.9.0.post0
PyYAML                       6.0.1
requests                     2.32.3
requests-oauthlib            2.0.0
rich                         13.7.1
rsa                          4.9
scipy                        1.11.4
setuptools                   59.6.0
six                          1.16.0
tensorboard                  2.11.2
tensorboard-data-server      0.6.1
tensorboard-plugin-wit       1.8.1
tensorflow                   2.11.1
tensorflow-datasets          4.4.0
tensorflow-estimator         2.11.0
tensorflow-io-gcs-filesystem 0.37.0
tensorflow-metadata          1.13.0
tensorstore                  0.1.61
termcolor                    2.4.0
toolz                        0.12.1
tqdm                         4.66.4
typing_extensions            4.12.2
urllib3                      2.2.1
Werkzeug                     3.0.3
wheel                        0.43.0
wrapt                        1.16.0
zipp                         3.19.2

python main.py --workdir=./imagenet --config=configs/default.py
2024-06-13 10:28:20.031363: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-06-13 10:28:20.612126: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2024-06-13 10:28:20.612204: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2024-06-13 10:28:20.612214: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Traceback (most recent call last):
  File "~/src/python/flax/examples/imagenet/main.py", line 29, in <module>
    import train
  File "~/src/python/flax/examples/imagenet/train.py", line 28, in <module>
    from flax import jax_utils
  File "~/src/python/flax/examples/imagenet/flax_imagenet_env/lib/python3.10/site-packages/flax/__init__.py", line 18, in <module>
    from .configurations import (
  File "~/src/python/flax/examples/imagenet/flax_imagenet_env/lib/python3.10/site-packages/flax/configurations.py", line 74, in <module>
    flax_filter_frames = define_bool_state(
  File "~/src/python/flax/examples/imagenet/flax_imagenet_env/lib/python3.10/site-packages/flax/configurations.py", line 40, in define_bool_state
    return jax_config.define_bool_state('flax_' + name, default, help)
AttributeError: 'Config' object has no attribute 'define_bool_state'

python -m venv mnist_env
source ./mnist_env/bin/activate
unset LD_LIBRARY_PATH
pip install "jax[cuda12]" flax tensorflow clu tensorflow-datasets
python main.py --workdir=/tmp/mnist --config=configs/default.py

absl-py==2.1.0
array_record==0.5.1
astunparse==1.6.3
cachetools==5.3.3
certifi==2024.6.2
charset-normalizer==3.3.2
chex==0.1.86
click==8.1.7
clu==0.0.12
contextlib2==21.6.0
dm-tree==0.1.8
docstring_parser==0.16
etils==1.7.0
flatbuffers==24.3.25
flax==0.8.4
fsspec==2024.6.0
gast==0.4.0
google-auth==2.30.0
google-auth-oauthlib==1.0.0
google-pasta==0.2.0
grpcio==1.64.1
h5py==3.11.0
idna==3.7
immutabledict==4.2.0
importlib_resources==6.4.0
jax==0.4.29
jax-cuda12-pjrt==0.4.29
jax-cuda12-plugin==0.4.29
jaxlib==0.4.29
keras==2.13.1
libclang==18.1.1
Markdown==3.6
markdown-it-py==3.0.0
MarkupSafe==2.1.5
mdurl==0.1.2
ml-collections==0.1.1
ml-dtypes==0.4.0
msgpack==1.0.8
nest-asyncio==1.6.0
numpy==1.24.3
nvidia-cublas-cu12==12.5.2.13
nvidia-cuda-cupti-cu12==12.5.39
nvidia-cuda-nvcc-cu12==12.5.40
nvidia-cuda-runtime-cu12==12.5.39
nvidia-cudnn-cu12==9.1.1.17
nvidia-cufft-cu12==11.2.3.18
nvidia-cusolver-cu12==11.6.2.40
nvidia-cusparse-cu12==12.4.1.24
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.5.40
oauthlib==3.2.2
opt-einsum==3.3.0
optax==0.2.2
orbax-checkpoint==0.5.16
packaging==24.1
promise==2.3
protobuf==3.20.3
psutil==5.9.8
pyarrow==16.1.0
pyasn1==0.6.0
pyasn1_modules==0.4.0
Pygments==2.18.0
PyYAML==6.0.1
requests==2.32.3
requests-oauthlib==2.0.0
rich==13.7.1
rsa==4.9
scipy==1.13.1
simple_parsing==0.1.5
six==1.16.0
tensorboard==2.13.0
tensorboard-data-server==0.7.2
tensorflow==2.13.1
tensorflow-datasets==4.9.6
tensorflow-estimator==2.13.0
tensorflow-io-gcs-filesystem==0.37.0
tensorflow-metadata==1.15.0
tensorstore==0.1.61
termcolor==2.4.0
toml==0.10.2
toolz==0.12.1
tqdm==4.66.4
typing_extensions==4.5.0
urllib3==2.2.1
Werkzeug==3.0.3
wrapt==1.16.0
zipp==3.19.2