Follow the step in readme. I got the following error
Traceback (most recent call last):
File "train.py", line 23, in <module>
from flax import jax_utils
File "/usr/local/lib/python3.8/dist-packages/flax/__init__.py", line 36, in <module>
from . import core
File "/usr/local/lib/python3.8/dist-packages/flax/core/__init__.py", line 15, in <module>
from .axes_scan import broadcast
File "/usr/local/lib/python3.8/dist-packages/flax/core/axes_scan.py", line 17, in <module>
import jax
File "/usr/local/lib/python3.8/dist-packages/jax/__init__.py", line 37, in <module>
from . import config as _config_module
File "/usr/local/lib/python3.8/dist-packages/jax/config.py", line 18, in <module>
from jax._src.config import config
File "/usr/local/lib/python3.8/dist-packages/jax/_src/config.py", line 27, in <module>
from jax import lib
File "/usr/local/lib/python3.8/dist-packages/jax/lib/__init__.py", line 73, in <module>
from jaxlib import xla_client
File "/usr/local/lib/python3.8/dist-packages/jaxlib/xla_client.py", line 27, in <module>
import ml_dtypes
File "/usr/local/lib/python3.8/dist-packages/ml_dtypes/__init__.py", line 29, in <module>
from ml_dtypes._custom_floats import bfloat16
ImportError: numpy.core.umath failed to import
Follow the step in readme. I got the following error