在Ubuntu上安装完JAX后，运行JAX程序报错：An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu。这个问题如何解决？

• 确保已安装CUDA和CuDNN。

• 安装CUDA和CuDNN

  • 安装CUDA：访问NVIDIA CUDA Toolkit Archive，下载并安装适合Ubuntu版本的CUDA。
  • 安装CuDNN：访问NVIDIA cuDNN，下载并安装与CUDA版本相匹配的CuDNN。

• 配置环境变量

  • 编辑bashrc：打开~/.bashrc文件，添加以下内容以配置CUDA环境变量。

    export CUDA_HOME=/usr/local/cuda-12.4
    export PATH=$PATH:$CUDA_HOME/bin
    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CUDA_HOME/lib64
    

  • 应用更改：运行source ~/.bashrc使更改生效。

• 安装CUDA-enabled jaxlib

  • 卸载旧版本：如果已安装旧版本的jax和jaxlib，使用pip uninstall jax jaxlib命令卸载。
  • 安装CUDA-enabled版本：根据系统CUDA和CuDNN版本，从JAX CUDA Releases下载并安装合适的jaxlib版本。例如，对于CUDA 12和CuDNN 8.9，使用命令：

    pip install jax==0.4.27
    pip install jaxlib-0.4.27+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl
    

• 验证GPU支持

  • 测试JAX：运行以下Python代码验证GPU是否被JAX正确识别。

    import jax
    print(jax.devices())
    

    如果输出显示cuda设备，说明安装成功。

以上步骤参考了多个来源，包括CSDN博客文章和官方安装指南。12567