linux下jax-GPU安装

安装命令

详细内容可查看jax官方文档

在已有cuda的情况下，先使用命令1，再使用命令2，即可完成安装
命令1：linux下 jax安装命令

pip install jax[cpu]==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_releases.html

将黑体部分，替换成想要的版本即可
命令2：linux下 jaxlib cuda安装命令
pip install jaxlib==0.3.25+cuda11.cudnn82 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
将黑体部分，替换成下面想要的版本即可
from versions: 0.1.52+cuda100, 0.1.52+cuda101, 0.1.52+cuda102, 0.1.52+cuda110, 0.1.55+cuda100, 0.1.55+cuda101, 0.1.55+cuda102, 0.1.55+cuda110, 0.1.56+cuda100, 0.1.56+cuda101, 0.1.56+cuda102, 0.1.56+cuda110, 0.1.56+cuda111, 0.1.57+cuda101, 0.1.57+cuda102, 0.1.57+cuda110, 0.1.57+cuda111, 0.1.58+cuda101, 0.1.58+cuda102, 0.1.58+cuda110, 0.1.58+cuda111, 0.1.59+cuda101, 0.1.59+cuda102, 0.1.59+cuda110, 0.1.59+cuda111, 0.1.60+cuda101, 0.1.60+cuda102, 0.1.60+cuda110, 0.1.60+cuda111, 0.1.61+cuda101, 0.1.61+cuda102, 0.1.61+cuda110, 0.1.61+cuda111, 0.1.61+cuda112, 0.1.62+cuda101, 0.1.62+cuda102, 0.1.62+cuda110, 0.1.62+cuda111, 0.1.62+cuda112, 0.1.63+cuda101, 0.1.63+cuda102, 0.1.63+cuda110, 0.1.63+cuda111, 0.1.63+cuda112, 0.1.64+cuda101, 0.1.64+cuda102, 0.1.64+cuda110, 0.1.64+cuda111, 0.1.64+cuda112, 0.1.65+cuda101, 0.1.65+cuda102, 0.1.65+cuda110, 0.1.65+cuda111, 0.1.65+cuda112, 0.1.66+cuda101, 0.1.66+cuda102, 0.1.66+cuda110, 0.1.66+cuda111, 0.1.67+cuda101, 0.1.67+cuda102, 0.1.67+cuda110, 0.1.67+cuda111, 0.1.68+cuda101, 0.1.68+cuda102, 0.1.68+cuda110, 0.1.68+cuda111, 0.1.69+cuda101, 0.1.69+cuda102, 0.1.69+cuda110, 0.1.69+cuda111, 0.1.70+cuda101, 0.1.70+cuda102, 0.1.70+cuda110, 0.1.70+cuda111, 0.1.71+cuda102, 0.1.71+cuda110, 0.1.71+cuda111, 0.1.72+cuda111, 0.1.73+cuda11.cudnn805, 0.1.73+cuda11.cudnn82, 0.1.74+cuda11.cudnn805, 0.1.74+cuda11.cudnn82, 0.1.75+cuda11.cudnn805, 0.1.75+cuda11.cudnn82, 0.1.76+cuda11.cudnn805, 0.1.76+cuda11.cudnn82, 0.3.0+cuda11.cudnn805, 0.3.0+cuda11.cudnn82, 0.3.2+cuda11.cudnn805, 0.3.2+cuda11.cudnn82, 0.3.5+cuda11.cudnn805, 0.3.5+cuda11.cudnn82, 0.3.7+cuda11.cudnn805, 0.3.7+cuda11.cudnn82, 0.3.8+cuda11.cudnn805, 0.3.8+cuda11.cudnn82, 0.3.10+cuda11.cudnn805, 0.3.10+cuda11.cudnn82, 0.3.14+cuda11.cudnn805, 0.3.14+cuda11.cudnn82, 0.3.15+cuda11.cudnn805, 0.3.15+cuda11.cudnn82, 0.3.18+cuda11.cudnn805, 0.3.18+cuda11.cudnn82, 0.3.20+cuda11.cudnn805, 0.3.20+cuda11.cudnn82, 0.3.22+cuda11.cudnn805, 0.3.22+cuda11.cudnn82, 0.3.24+cuda11.cudnn805, 0.3.24+cuda11.cudnn82, 0.3.25+cuda11.cudnn805, 0.3.25+cuda11.cudnn82

查看是否安装成功

在命令行中输入

python

进入后输入

import jax
print(jax.devices())