Revision 4d8bce1b8515c68a8432d379d964ea3e6bea463c authored by Yash Katariya on 19 October 2021, 02:55:51 UTC, committed by jax authors on 19 October 2021, 02:56:22 UTC
```
# Installs Cuda 11 with Cudnn 8.2
$ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
```

PiperOrigin-RevId: 404134291
1 parent 6bd0c72
History
File Mode Size
BUILD.bazel -rw-r--r-- 1.7 KB
Dockerfile -rw-r--r-- 1.5 KB
LICENSE.txt -rw-r--r-- 225.2 KB
build.py -rwxr-xr-x 17.2 KB
build_jaxlib_wheels.sh -rwxr-xr-x 311 bytes
build_jaxlib_wheels_helpers.sh -rw-r--r-- 935 bytes
build_jaxlib_wheels_macos.sh -rwxr-xr-x 1.5 KB
build_wheel.py -rw-r--r-- 9.8 KB
build_wheel_docker_entrypoint.sh -rwxr-xr-x 1.4 KB
install_cuda.sh -rw-r--r-- 1.6 KB
test-requirements.txt -rw-r--r-- 83 bytes

back to top