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

Readd cuda12_pip extra to keep CI users happy.

Browse Source
This commit is contained in:
Peter Hawkins 2024-06-13 12:36:54 -04:00
parent 2679ece82d
commit dcb7b3c3f1
1 changed files with 7 additions and 0 deletions

7
setup.py

@ -95,6 +95,13 @@ setup(
f"jax-cuda12-plugin[with_cuda]>={_current_jaxlib_version},<={_jax_version}",
],
# Deprecated alias for cuda12, kept to avoid breaking users who wrote
# cuda12_pip in their CI.
'cuda12_pip': [
f"jaxlib=={_current_jaxlib_version}",
f"jax-cuda12-plugin[with_cuda]>={_current_jaxlib_version},<={_jax_version}",
],
# Target that does not depend on the CUDA pip wheels, for those who want
# to use a preinstalled CUDA.
'cuda12_local': [