无法安装特定的 JAX jaxlib GPU 版本答案

【问题标题】：Unable to Install Specific JAX jaxlib GPU version无法安装特定的 JAX jaxlib GPU 版本
【发布时间】：2020-12-29 01:58:10
【问题描述】：

我正在尝试安装特定版本的 jaxlib 以使用我的 CUDA 和 cuDNN 版本。按照自述文件，我正在尝试

pip install --upgrade jax jaxlib==0.1.52+cuda101 -f https://storage.googleapis.com/jax-releases/jax_releases.html

这会返回以下错误：

ERROR: Requested jaxlib==0.1.52+cuda101 from https://storage.googleapis.com/jax-releases/cuda101/jaxlib-0.1.52%2Bcuda101-cp37-none-manylinux2010_x86_64.whl has different version in metadata: '0.1.52'

有谁知道是什么原因造成的或者如何解决这个错误？

【问题讨论】：

标签： installation gpu failed-installation jax

【解决方案1】：

此错误似乎来自 pip 20.3.X 及更高版本中的新检查，可能与新的依赖关系解析器有关。我可以使用 pip 版本 20.3.3 重现此错误，但使用 pip 版本 20.2.4 可以正确安装包。

最简单的方法可能是首先降级 pip；即

pip install pip==20.2.4

然后继续您的 jaxlib 安装。

【讨论】：

你会推荐 conda-forge 作为安装 JAX 的好方法吗？
对于 CPU，这可能有效。但我不认为 conda-forge 目前提供与 GPU 兼容的 jaxlib 构建。

【解决方案2】：

请注意，jax 和 jaxlib 的版本必须匹配。你可以使用类似的东西：

$ pip install --upgrade jax==0.3.2 jaxlib==0.3.2+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_releases.html

另一种解决方法是首先选择specific version of jax and jaxlib from the available wheel files and then install those。

$ pip install https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.1.76+cuda11.cudnn82-cp39-none-manylinux2010_x86_64.whl

【讨论】：