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 install instructions for AMD (ROCm) GPUs to point to AMD's docker container.

Browse Source
This commit is contained in:
Peter Hawkins 2023-09-26 09:49:06 -04:00
parent bfc3193da0
commit fb97f48870
2 changed files with 11 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
README.md
View File

@ -385,7 +385,7 @@ Some standouts:
| CPU | yes | yes (build from source) | yes | yes | yes | yes |
| NVIDIA GPU | yes | yes (build from source) | no | n/a | no | experimental |
| Google TPU | yes | n/a | n/a | n/a | n/a | n/a |
| AMD GPU | experimental (build from source) | no | no | n/a | no | no |
| AMD GPU | experimental | no | no | n/a | no | no |
| Apple GPU | n/a | no | experimental | experimental | n/a | n/a |
@ -396,11 +396,11 @@ Some standouts:
| CPU | `pip install -U jax[cpu]` |
| NVIDIA GPU | `pip install -U "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html` |
| Google TPU | `pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html` |
| AMD GPU | [Build from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus). |
| AMD GPU | Use [Docker](https://hub.docker.com/r/rocm/jax) or [build from source](https://jax.readthedocs.io/en/latest/developer.html#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus). |
| Apple GPU | Follow [Apple's instructions](https://developer.apple.com/metal/jax/). |
See [the documentation](https://jax.readthedocs.io/en/latest/installation.html)
for information on alternative installation strategies. These includes compiling
for information on alternative installation strategies. These include compiling
from source, installing with Docker, using other versions of CUDA, a
community-supported conda build, and answers to some frequently-asked questions.

9
docs/installation.md
View File

@ -13,7 +13,7 @@ source](developer.md#building-from-source).
| CPU | [yes](#cpu) | [yes (build from source)](developer.md#building-from-source) | [yes](#cpu) | [yes](#cpu) | [yes](#cpu) | [yes](#cpu) |
| NVIDIA GPU | [yes](#nvidia-gpu) | [yes (build from source)](developer.md#building-from-source) | no | n/a | no | [experimental](#nvidia-gpu) |
| Google TPU | [yes](#google-tpu) | n/a | n/a | n/a | n/a | n/a |
| AMD GPU | [experimental (build from source)](developer.md#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus) | no | no | n/a | no | no |
| AMD GPU | [experimental](#amd-gpu) | no | no | n/a | no | no |
| Apple GPU | n/a | no | [experimental](#apple-gpu) | [experimental](#apple-gpu) | n/a | n/a |
@ -194,6 +194,13 @@ There are several caveats with the Metal plugin:
`jaxlib`. This restriction will be relaxed over time as the plugin API
matures.
## AMD GPU
JAX has experimental ROCM support. There are two ways to install JAX:
* use [AMD's docker container](https://hub.docker.com/r/rocm/jax), or
* [build from source](developer.md#additional-notes-for-building-a-rocm-jaxlib-for-amd-gpus).
## Conda
### Conda installation