jax开源库安装问题

介绍

jax是谷歌开发的一个使用GPU、TPU加速一些结算的一个开源项目，部分中文介绍在一些文章上可以了解，https://zhuanlan.zhihu.com/p/468223891，可在GitHub上找到git clone链接

具体安装

一、使用pip安装

pip安装的方法很简单，但也容易出现问题，比如由于下载的pip版本或者其他package的版本会影响安装，有可能只下载了jax但没有下载jaxlib，但jax需要jaxlib的支持，所以后续出现问题很容易找不到原因。
安装之前，最好先更新pip（如果是最新版本请忽略），且最好不要用pip install --upgrade pip，使用conda更新更方便些，打开cmd要使用管理员权限
'''
//cpu版本
pip install --upgrade "jax[cpu]" 
//gpu版本
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.htm
//TPU版本
pip install "jax[tpu]>=0.2.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
'''

二、使用GitBash安装

首先需要去GitHub官网下载gitbash，然后安装，或者看博文：https://blog.csdn.net/aixueai/article/details/119574135
找到Python的包存放目录，一般在InstallLocation\Lib\site-packages中，右键，可以看到GitBash here

‘’’
//在GitBash界面内输入指令
git init
git clone https://github.com/google/jax.git
//等待clone完毕
‘’’

最后打开下载好的jax文件，在里面可以看到setup.py文件，在这个目录下打开cmd

'''
python setup.py install 
//等待安装，再返回site-package目录，这时候可以看到jaxlib在里面了
//这时候jax安装好了
'''