该项目目前仍然没有官方的Windows支持,需要自己编译。
安装说明
该库安装时分为两部分:
- jaxlib,该库平台相关,目前没有官方的编译
- jax,该库依赖jaxlib,平台无关,可以直接安装。
找到目前一个还活跃的jaxlib非官方编译服务:
https://github.com/cloudhan/jax-windows-builder
可以使用类似
pip install jaxlib[cuda111] -f https://whls.blob.core.windows.net/unstable/index.html --use-deprecated legacy-resolver
的方式直接安装,不同版本详见项目主页。