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

Update changelog and install doc to mention cuda 12.3 switch

Browse Source 
PiperOrigin-RevId: 605473026
This commit is contained in:
Yash Katariya 2024-02-08 17:21:16 -08:00 committed by jax authors
parent 9b320f23f0
commit 73e3dedf9d
2 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
CHANGELOG.md
View File

@ -127,6 +127,8 @@ Remember to align the itemized text with the first line of an item within a list
* Changes
* JAX now supports CUDA 12.3 and CUDA 11.8. Support for CUDA 12.2 has been
dropped.
* `cost_analysis` now works with cross-compiled `Compiled` objects (i.e. when
using `.lower().compile()` with a topology object, e.g., to compile for
Cloud TPU from a non-TPU computer).

2
docs/installation.md
View File

@ -114,7 +114,7 @@ able to use the
that NVIDIA provides for this purpose.
JAX currently ships two CUDA wheel variants:
* CUDA 12.2, cuDNN 8.9, NCCL 2.16
* CUDA 12.3, cuDNN 8.9, NCCL 2.16
* CUDA 11.8, cuDNN 8.6, NCCL 2.16
You may use a JAX wheel provided the major version of your CUDA, cuDNN, and NCCL