diff --git a/.bazelrc b/.bazelrc index 9c4639ed8..30891e0bf 100644 --- a/.bazelrc +++ b/.bazelrc @@ -66,7 +66,7 @@ build:mkl_open_source_only --define=tensorflow_mkldnn_contraction_kernel=1 build:cuda --repo_env TF_NEED_CUDA=1 # "sm" means we emit only cubin, which is forward compatible within a GPU generation. # "compute" means we emit both cubin and PTX, which is larger but also forward compatible to future GPU generations. -build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="sm_52,sm_60,sm_70,compute_80" +build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="sm_52,sm_60,sm_70,sm_80,compute_90" build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain build:cuda --@local_config_cuda//:enable_cuda build:cuda --@xla//xla/python:enable_gpu=true diff --git a/CHANGELOG.md b/CHANGELOG.md index 1774e2382..b0ea254d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,11 @@ Remember to align the itemized text with the first line of an item within a list ## jaxlib 0.4.12 +* Changes + * Include PTX/SASS for Hopper (SM version 9.0+) GPUs. Previous + versions of jaxlib should work on Hopper but would have a long + JIT-compilation delay the first time a JAX operation was executed. + ## jax 0.4.11 (May 31, 2023) * Deprecations