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

Add Cuda 11.4 with cudnn 8.2 and cudnn 8.0.5 release builds

Browse Source 
PiperOrigin-RevId: 403661187
This commit is contained in:
Yash Katariya 2021-10-16 16:13:15 -07:00 committed by jax authors
parent 267a4ca4cb
commit e6e81ba885
2 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.bazelrc
View File

@ -210,7 +210,6 @@ build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --extra_execution_platforms="@ubu
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_platform//:platform"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_cuda"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_tensorrt"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_nccl"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.7 --config=rbe_linux_cuda11.4_cudnn8.0.5_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.4-cudnn8.0.5-tensorrt7.2_config_python3.7"
build:rbe_linux_cuda11.4_cudnn8.0.5_nvcc_py3.7 --python_path="/usr/local/bin/python3.7"

5
jaxlib/setup.py
View File

@ -21,8 +21,9 @@ with open('jaxlib/version.py') as f:
exec(f.read(), globals())
cuda_version = os.environ.get("JAX_CUDA_VERSION")
if cuda_version:
__version__ += "+cuda" + cuda_version.replace(".", "")
cudnn_version = os.environ.get("JAX_CUDNN_VERSION")
if cuda_version and cudnn_version:
__version__ += f"+cuda{cuda_version.replace('.', '')}_cudnn{cudnn_version.replace('.', '')}"
setup(
name='jaxlib',