mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 21:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add jax[_ci] option to account for the new release process.

Browse Source 
PiperOrigin-RevId: 427802081
This commit is contained in:
Yash Katariya 2022-02-10 11:29:06 -08:00 committed by jax authors
parent 0b135fbfe8
commit d82bcc2a0c
1 changed files with 5 additions and 1 deletions

6
setup.py

@ -14,8 +14,9 @@
from setuptools import setup, find_packages
# The following should be updated with each new jaxlib release.
_current_jaxlib_version = '0.1.76'
# The following should be updated with each new jaxlib release.
_latest_jaxlib_version_on_pypi = '0.1.76'
_available_cuda_versions = ['11']
_default_cuda_version = '11'
_available_cudnn_versions = ['82', '805']
@ -52,6 +53,9 @@ setup(
# $ pip install jax[cpu]
'cpu': [f'jaxlib=={_current_jaxlib_version}'],
# Used only for CI builds that install JAX from github HEAD.
'_ci': [f'jaxlib=={_latest_jaxlib_version_on_pypi}'],
# Cloud TPU VM jaxlib can be installed via:
# $ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_releases.html
'tpu': [f'jaxlib=={_current_jaxlib_version}',