jax和jaxlib的cuda版本安装

见贤思齐

笔者花费时间才在Ubuntu20.04适配上jax和jaxlib的cuda版本安装，以及chex版本。版本展示本人版本展示jax==0.4.27，jaxlib==0.4.27+cuda12.cudnn89，chex==0.1.86。安装过程cuda以及环境变量配置过程首先安装cuda12.4和cudnn8.9，给出链接：cuda：https://developer.nvidia.com/cuda-toolkit-archivecudnn：https://developer.nvidia.com/rdp/cudnn-download安装完成cuda和cudnn后，务必将cuda写入环境变量，给出模板示例，特别是CUPTI库文件：vim~/.bashrc#cuda安装位置exportCUDA_HOME=/usr/local/cuda-12.4exportPATH=$PATH:/usr/local/cuda-12.4/binexportLD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.4/lib64exportLD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-12.4/extras/CUPTI/lib64exportLD_LIBRARY_PATH=/usr/local/cuda-12.4/targets/x86_64-linux/libLD_LIBRARY_PATHsource~/.bashrc给出笔者配置bashrc图片参考:jax和jaxlib安装过程此处首先给出官方cuda版本whl文件下载地址：jaxlib：storage.googleapis.com/jax-releases/jax_cuda_releases.html选择cuda12.cudnn89对应自己需要版本，根据自己要求选择，笔者选择为0.4.27，则执行：pipinstalljax==0.4.27#此处jaxlib笔者为whl本地安装pipinstalljaxlib-0.4.27+cuda12.cudnn89-cp39-cp39-manylinux2014_x86_64.whl切记笔者chex==0.1.86，chex对jax也有要求，容易遇到bug，全是泪与恨。验证是否GPU可用在终端内验证：>>>importjax>>>print(jax.devices())[cuda(id=0)]正确输出cudaid即为成功。参考（Reference）JAX:库安装和GPU使用，解决不能识别gpu问题_jax找不到gpu-CSDN博客JAX安装过程中遇到的各种坑-CSDN博客