mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix jax releases URL in setup.py comments

Browse Source
This commit is contained in:
Jake VanderPlas 2022-10-25 10:39:39 -07:00
parent adcb0f58e8
commit 41b815bf33
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
setup.py
View File

@ -92,12 +92,12 @@ setup(
# CUDA installations require adding jax releases URL; e.g.
# Cuda installation defaulting to a CUDA and Cudnn version defined above.
# $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html
# $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
'cuda': [f"jaxlib=={_current_jaxlib_version}+cuda{_default_cuda_version}.cudnn{_default_cudnn_version}"],
# CUDA installations require adding jax releases URL; e.g.
# $ pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html
# $ pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
# $ pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
# $ pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
**{f'cuda{cuda_version}_cudnn{cudnn_version}': f"jaxlib=={_current_jaxlib_version}+cuda{cuda_version}.cudnn{cudnn_version}"
for cuda_version in _available_cuda_versions for cudnn_version in _available_cudnn_versions}
},