flax与jax、jaxlib的安装顺序和版本匹配

---

前言

在复现代码时需要用到这三个库，在安装的时候遇到了一些问题，现将解决措施记录如下。

---

一、安装顺序及方法

顺序是：jaxlib→jax→flax，安装之后记得检查能否使用GPU
jaxlib安装方法：
（1）Windows：https://whls.blob.core.windows.net/unstable/index.html
（2）Linux或Mac：直接在https://pypi.org/上下载安装包并安装
jax安装：https://blog.csdn.net/Papageno_Xue/article/details/125754893
flax也是直接在pypi上下载安装包并安装。

二、版本匹配

1.基础环境

这三个库是在TensorFlow环境当中安装的，因此版本首先需要与tf版本匹配。基础环境的版本如下：
ml-collections 0.1.1
tensorflow-gan 2.1.0
tensorflow-gpu 2.9.0
tensorflow-hub 0.12.0
tensorflow-io 0.28.0
tensorflow-addons 0.17.1
tensorflow-datasets 4.3.0
tensorflow-probability 0.16.0
tensorboard 2.9.1
cuda 11.2，cudnn 8.2.0

---

2.三个库的版本匹配

jaxlib 0.3.14+cuda11.cudnn82；jax 0.3.14；flax 0.5.3。
这里介绍一个很有用的工具pipdeptree，如果需要查看某个库的依赖库有哪些，包括对应的版本，可以用 pipdeptree -p 库名 来查看。

总结

本文介绍了flax与jax、jaxlib的安装顺序和版本匹配，希望对大家有帮助。