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
.github
benchmarks
build
cloud_tpu_colabs
design_notes
docs
examples
images
jax
jaxlib
tests
third_party
.bazelrc -rw-r--r-- 15.1 KB
.bazelversion -rw-r--r-- 6 bytes
.gitignore -rw-r--r-- 291 bytes
.pre-commit-config.yaml -rw-r--r-- 563 bytes
.readthedocs.yml -rw-r--r-- 528 bytes
CHANGELOG.md -rw-r--r-- 42.8 KB
CONTRIBUTING.md -rw-r--r-- 150 bytes
LICENSE -rw-r--r-- 11.1 KB
README.md -rw-r--r-- 22.6 KB
WORKSPACE -rw-r--r-- 1.3 KB
conftest.py -rw-r--r-- 852 bytes
mypy.ini -rw-r--r-- 523 bytes
pylintrc -rw-r--r-- 1.6 KB
pytest.ini -rw-r--r-- 768 bytes
setup.cfg -rw-r--r-- 329 bytes
setup.py -rw-r--r-- 3.2 KB

README.md

back to top