如何安装jax？

[MFCadq]

在运行gpu版本的安装命令之后，出现无法读取到config文件的问题。。
ModuleNotFoundError: No module named 'jax.config'
请问一下如何解决呢？