关于jax，tensorflow的安装详细教程

jax在下面pycharm中有提示你在哪里下载，所以如果下次遇到有什么要下载的包都可以尝试这种方式去下载（在找找pycharm有没有提示，然后直接点网站看教程，这嘎是官网的教程，很详细）

 为了下载一个版本只有JAX的cpu，可能要用到当地的开发，你可以用.....(我英语很差的！)

依次输入以下代码

pip install --upgrade pip
pip install --upgrade "jax[cpu]"

pip install --upgrade pip
# Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html  # Note: wheels only available on linux.

你可以为jaxlib显式指定一个特定的CUDA和cuDNN版本:  

pip install --upgrade pip

# Installs the wheel compatible with Cuda 11 and cudnn 8.2 or newer.
pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

# Installs the wheel compatible with Cuda 11 and cudnn 8.0.5 or newer.
pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html

你可以通过命令找到你的CUDA版本:  

nvcc --version

就可以了，还有什么不会的下面提问