Skip to content

Commit 11eeff0

Browse files
Google-ML-AutomationGoogle-ML-Automation
committed
Merge pull request jax-ml#22410 from garymm:patch-1
PiperOrigin-RevId: 688265373
2 parents 5de8780 + dc908b4 commit 11eeff0

File tree

2 files changed

+22
-22
lines changed

2 files changed

+22
-22
lines changed

‎README.md‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -383,24 +383,24 @@ Some standouts:
383383

384384
### Supported platforms
385385

386-
| | Linux x86_64 | Linux aarch64 | Mac x86_64 | Mac ARM | Windows x86_64 | Windows WSL2 x86_64 |
386+
| | Linux x86_64 | Linux aarch64 | Mac x86_64 | Mac aarch64 | Windows x86_64 | Windows WSL2 x86_64 |
387387
|------------|--------------|---------------|--------------|--------------|----------------|---------------------|
388388
| CPU | yes | yes | yes | yes | yes | yes |
389389
| NVIDIA GPU | yes | yes | no | n/a | no | experimental |
390390
| Google TPU | yes | n/a | n/a | n/a | n/a | n/a |
391-
| AMD GPU | yes | no | no | n/a | no | no |
392-
| Apple GPU | n/a | no | experimental | experimental | n/a | n/a |
391+
| AMD GPU | yes | no | experimental | n/a | no | no |
392+
| Apple GPU | n/a | no | n/a | experimental | n/a | n/a |
393393

394394

395395
### Instructions
396396

397-
| Hardware | Instructions |
398-
|------------|-----------------------------------------------------------------------------------------------------------------|
399-
| CPU | `pip install -U jax` |
400-
| NVIDIA GPU | `pip install -U "jax[cuda12]"` |
401-
| Google TPU | `pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html` |
402-
| AMD GPU | Use [Docker](https://hub.docker.com/r/rocm/jax-community/tags), [pre-built wheels](https://github.com/ROCm/jax/releases), or [build from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus). |
403-
| Apple GPU | Follow [Apple's instructions](https://developer.apple.com/metal/jax/). |
397+
| Platform | Instructions |
398+
|-----------------|-----------------------------------------------------------------------------------------------------------------|
399+
| CPU | `pip install -U jax` |
400+
| NVIDIA GPU | `pip install -U "jax[cuda12]"` |
401+
| Google TPU | `pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html` |
402+
| AMD GPU (Linux) | Use [Docker](https://hub.docker.com/r/rocm/jax-community/tags), [pre-built wheels](https://github.com/ROCm/jax/releases), or [build from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus). |
403+
| Mac GPU | Follow [Apple's instructions](https://developer.apple.com/metal/jax/). |
404404

405405
See [the documentation](https://jax.readthedocs.io/en/latest/installation.html)
406406
for information on alternative installation strategies. These include compiling

‎docs/installation.md‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,13 @@ different builds for different operating systems and accelerators.
2828

2929
The table below shows all supported platforms and installation options. Check if your setup is supported; and if it says _"yes"_ or _"experimental"_, then click on the corresponding link to learn how to install JAX in greater detail.
3030

31-
| | Linux, x86_64 | Linux, aarch64 | macOS, Intel x86_64, AMD GPU | macOS, Apple Silicon, ARM-based | Windows, x86_64 | Windows WSL2, x86_64 |
32-
|------------------|---------------------------------------|--------------------------------|----------------------------------------|----------------------------------------|-------------------------|-----------------------------------------|
33-
| CPU | {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>`| {ref}`yes <install-cpu>`| {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>`|
34-
| NVIDIA GPU | {ref}`yes <install-nvidia-gpu>` | {ref}`yes <install-nvidia-gpu>` | no | n/a | no | {ref}`experimental <install-nvidia-gpu>` |
35-
| Google Cloud TPU | {ref}`yes <install-google-tpu>` | n/a | n/a | n/a | n/a | n/a |
36-
| AMD GPU | {ref}`experimental <install-amd-gpu>` | no | no | n/a | no | no |
37-
| Apple GPU | n/a | no | {ref}`experimental <install-apple-gpu>` | {ref}`experimental <install-apple-gpu>` | n/a | n/a |
31+
| | Linux, x86_64 | Linux, aarch64 | Mac, x86_64 | Mac, aarch64 | Windows, x86_64 | Windows WSL2, x86_64 |
32+
|------------------|---------------------------------------|---------------------------------|---------------------------------------|---------------------------------------|--------------------------|------------------------------------------|
33+
| CPU | {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>` | {ref}`yes <install-cpu>` |
34+
| NVIDIA GPU | {ref}`yes <install-nvidia-gpu>` | {ref}`yes <install-nvidia-gpu>` | no | n/a | no | {ref}`experimental <install-nvidia-gpu>` |
35+
| Google Cloud TPU | {ref}`yes <install-google-tpu>` | n/a | n/a | n/a | n/a | n/a |
36+
| AMD GPU | {ref}`experimental <install-amd-gpu>` | no | {ref}`experimental <install-mac-gpu>` | n/a | no | no |
37+
| Apple GPU | n/a | no | n/a | {ref}`experimental <install-mac-gpu>` | n/a | n/a |
3838

3939

4040
(install-cpu)=
@@ -204,12 +204,12 @@ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_relea
204204
For users of Colab (https://colab.research.google.com/), be sure you are
205205
using *TPU v2* and not the older, deprecated TPU runtime.
206206

207-
(install-apple-gpu)=
208-
## Apple Silicon GPU (ARM-based)
207+
(install-mac-gpu)=
208+
## Mac GPU
209209

210-
### pip installation: Apple ARM-based Silicon GPUs
210+
### pip installation
211211

212-
Apple provides an experimental Metal plugin for Apple ARM-based GPU hardware. For details,
212+
Apple provides an experimental Metal plugin. For details,
213213
refer to
214214
[Apple's JAX on Metal documentation](https://developer.apple.com/metal/jax/).
215215

@@ -223,7 +223,7 @@ refer to
223223
matures.
224224

225225
(install-amd-gpu)=
226-
## AMD GPU
226+
## AMD GPU (Linux)
227227

228228
JAX has experimental ROCm support. There are two ways to install JAX:
229229

0 commit comments

Comments
 (0)