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

Update README discussion of Windows support:

Browse Source 
* right next to the pip installation instructions, mention they don't work for Windows;
* add a link to #5795 for an unofficial discussion of Windows native support
This commit is contained in:
Matthew Johnson 2022-07-19 15:04:16 -07:00
parent 6ea9e4d4dd
commit 3ab8637a73
1 changed files with 19 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
README.md
View File

@ -388,15 +388,19 @@ Some standouts:
JAX is written in pure Python, but it depends on XLA, which needs to be JAX is written in pure Python, but it depends on XLA, which needs to be
installed as the `jaxlib` package. Use the following instructions to install a installed as the `jaxlib` package. Use the following instructions to install a
binary package with `pip`, or to build JAX from source. binary package with `pip`, or to [build JAX from
source](https://jax.readthedocs.io/en/latest/developer.html#building-from-source).
We support installing or building `jaxlib` on Linux (Ubuntu 16.04 or later) and We support installing or building `jaxlib` on Linux (Ubuntu 16.04 or later) and
macOS (10.12 or later) platforms. Windows users can use JAX on CPU and GPU via macOS (10.12 or later) platforms.
the
[Windows Subsystem for Linux](https://docs.microsoft.com/en-us/windows/wsl/about). Windows users can use JAX on CPU and GPU via the [Windows Subsystem for
There is some initial native Windows support, but since it is still somewhat Linux](https://docs.microsoft.com/en-us/windows/wsl/about). In addition, there
immature, there are no binary releases and it must be is some initial community-driven native Windows support, but since it is still
[built from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-jaxlib-from-source-on-windows). somewhat immature, there are no official binary releases and it must be [built
from source for Windows](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-jaxlib-from-source-on-windows).
For an unofficial discussion of native Windows builds, see also the [Issue #5795
thread](https://github.com/google/jax/issues/5795).
### pip installation: CPU ### pip installation: CPU
@ -410,6 +414,8 @@ pip install --upgrade "jax[cpu]"
On Linux, it is often necessary to first update `pip` to a version that supports On Linux, it is often necessary to first update `pip` to a version that supports
`manylinux2014` wheels. `manylinux2014` wheels.
**These `pip` installations do not work with Windows, and may fail silently; see
[above](#installation).**
### pip installation: GPU (CUDA) ### pip installation: GPU (CUDA)
@ -422,12 +428,12 @@ package.
JAX provides pre-built CUDA-compatible wheels for **Linux only**, JAX provides pre-built CUDA-compatible wheels for **Linux only**,
with CUDA 11.1 or newer, and CuDNN 8.0.5 or newer. Other combinations of with CUDA 11.1 or newer, and CuDNN 8.0.5 or newer. Other combinations of
operating system, CUDA, and CuDNN are possible, but require building from operating system, CUDA, and CuDNN are possible, but require [building from
source. source](https://jax.readthedocs.io/en/latest/developer.html#building-from-source).
* CUDA 11.1 or newer is *required*. * CUDA 11.1 or newer is *required*.
* You may be able to use older CUDA versions if you build from source, but * You may be able to use older CUDA versions if you [build from source](https://jax.readthedocs.io/en/latest/developer.html#building-from-source),
there are known bugs in CUDA in all CUDA versions older than 11.1, so we but there are known bugs in CUDA in all CUDA versions older than 11.1, so we
do not ship prebuilt binaries for older CUDA versions. do not ship prebuilt binaries for older CUDA versions.
* The supported cuDNN versions for the prebuilt wheels are: * The supported cuDNN versions for the prebuilt wheels are:
* cuDNN 8.2 or newer. We recommend using the cuDNN 8.2 wheel if your cuDNN * cuDNN 8.2 or newer. We recommend using the cuDNN 8.2 wheel if your cuDNN
@ -455,6 +461,8 @@ pip install --upgrade pip
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
``` ```
**These `pip` installations do not work with Windows, and may fail silently; see
[above](#installation).**
The jaxlib version must correspond to the version of the existing CUDA The jaxlib version must correspond to the version of the existing CUDA
installation you want to use. You can specify a particular CUDA and CuDNN installation you want to use. You can specify a particular CUDA and CuDNN