jax安装踩坑(1) ImportError: cannot import name ‘linear_util‘ from ‘jax‘

https://stackoverflow.com/questions/78210393/cannot-import-name-linear-util-from-jax

To fix your issue, you’ll either need to install an older version of JAX which still has jax.linear_util, or update to a newer version of flax which is compatible with more recent JAX versions.

评论区说，要么降级jax，要么升级flax。笔者这边升级flax：

pip install -U flax

通常来说问题到这里就解决了。

作者的特殊情况

但是flax与jax、jaxlib有版本兼容，因此安装flax会同步替换jax的版本。而在笔者情况下，jax不能安装过高版本的flax，最后选择了如下版本：

pip install flax==0.5.3

然后再安装：

pip install jax[cuda12]==0.4.30