在linux上使用jax的gpu加速

版本

在命令窗口输入nvidia-smi查看

NVIDIA-SMI 520.61.05    Driver Version: 520.61.05    CUDA Version: 11.8 

找到对应的版本https://jax.readthedocs.io/en/latest/installation.html

pip install --upgrade pip

# CUDA 12 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# CUDA 11 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

那么就是

pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

等待安装

检查是否安装成功

import jax  
import jax.numpy as jnp  
  
print("Number of GPUs available:", jax.local_device_count())  
  
x = jnp.ones((2, 2))  
y = jax.lax.scan(lambda carry, x: (carry + x, x), jnp.zeros(2), x)  

y_val = jax.device_put(y, jax.devices('gpu')[0])  
print(y_val)

如果没有报错并且出现结果即可