jax 不识别gpu

以下内容在 x86 linux cudagpu 测试通过

错误信息

jax-0.4.25 jaxlib-0.4.25 不识别gpu: RuntimeError: Unknown backend: ‘gpu’ requested, but no platforms that are instances of gpu are present. Platforms present are: cpu

最终解决方案

根据cudnn版本和py版本去 https://storage.googleapis.com/jax-releases/jax_cuda_releases.html 查对于的jaxlib版本来进行安装。

倒腾过程

pip install --upgrade “jax[cuda]” 会装jax-0.2.22 jaxlib-0.4.25

jax-0.2.22：引入就报错。ImportError: cannot import name ‘pocketfft’ from ‘jaxlib’

pip uninstall jax jaxlib

https://pypi.org/project/jax/

pip install -U “jax[cuda12_pip]” -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

装完就是 jax-0.4.25 jaxlib-0.4.25+cuda12.cudnn89

jax.devices() 会爆下面的错误

CUDA backend failed to initialize: Found cuDNN version 8400, but JAX was built against version 8906, which is newer. The copy of cuDNN that is installed must be at least as new as the version against which JAX was built. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

https://blog.csdn.net/Papageno_Xue/article/details/125754893

不存在 /usr/local/cuda-12.1/include/cudnn_version.h,/usr/local/cuda/include/cudnn_version.h 显示是840，和上面提示的错误信息一直

这里确定了cudaa 版本8400,然后我是py39

jaxlib version

https://storage.googleapis.com/jax-releases/jax_cuda_releases.html 里没有 0.4.25+84的，cudnn84以下有的82的，但是是cuda11，直接 pip install -U “jax[cuda11_pip]” -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html 的话装的是jaxlib==0.4.25+cuda11.cudnn86，根据 py39+cudaa82 查看得知最高 jaxlib0.4.7

pip install jax0.4.7 jaxlib0.4.7+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

    >>> import jax
    >>> jax.devices()
    [StreamExecutorGpuDevice(id=0, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=1, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=2, process_index=0, slice_index=0), StreamExecutorGpuDevice(id=3, process_index=0, slice_index=0)]