ImportError: numpy.core.umath failed to import

[HedgeHao]

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