Skip to content

Releases: ROCm/jax

rocm-jax-v0.4.35

05 Dec 03:51
c47e0f2
Compare
Choose a tag to compare

ROCm JAX 0.4.35 Release

ROCm 6.0.3, 6.1.3 and 6.2.4

Python 3.10.15

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jaxlib-0.4.35-cp310-cp310-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jax_rocm60_pjrt-0.4.35-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jax_rocm60_plugin-0.4.35-cp310-cp310-manylinux_2_28_x86_64.whl

Python 3.11.10

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jaxlib-0.4.35-cp311-cp311-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jax_rocm60_pjrt-0.4.35-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jax_rocm60_plugin-0.4.35-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12.7

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jaxlib-0.4.35-cp312-cp312-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jax_rocm60_pjrt-0.4.35-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.35/jax_rocm60_plugin-0.4.35-cp312-cp312-manylinux_2_28_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jax-v0.4.35.tar.gz

Docker Images

ROCm 6.0.3

docker pull rocm/jax-community:rocm6.0.3-jax0.4.35-py3.10.15
docker pull rocm/jax-community:rocm6.0.3-jax0.4.35-py3.11.10
docker pull rocm/jax-community:rocm6.0.3-jax0.4.35-py3.12.7

ROCm 6.1.3

docker pull rocm/jax-community:rocm6.1.3-jax0.4.35-py3.10.15
docker pull rocm/jax-community:rocm6.1.3-jax0.4.35-py3.11.10
docker pull rocm/jax-community:rocm6.1.3-jax0.4.35-py3.12.7

ROCm 6.2.4

docker pull rocm/jax-community:rocm6.2.4-jax0.4.35-py3.10.15
docker pull rocm/jax-community:latest-py310
docker pull rocm/jax-community:rocm6.2.4-jax0.4.35-py3.11.10
docker pull rocm/jax-community:latest-py311
docker pull rocm/jax-community:rocm6.2.4-jax0.4.35-py3.12.7
docker pull rocm/jax-community:latest-py312
docker pull rocm/jax-community:latest

rocm-jax-v0.4.34

25 Oct 15:58
Compare
Choose a tag to compare

ROCm JAX 0.4.34 Release

ROCm 6.0.3, 6.1.3 and 6.2.3

Python 3.10.15

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jaxlib-0.4.34-cp310-cp310-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jax_rocm60_pjrt-0.4.34-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jax_rocm60_plugin-0.4.34-cp310-cp310-manylinux_2_28_x86_64.whl

Python 3.11.10

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jaxlib-0.4.34-cp311-cp311-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jax_rocm60_pjrt-0.4.34-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jax_rocm60_plugin-0.4.34-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12.6

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jaxlib-0.4.34-cp312-cp312-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jax_rocm60_pjrt-0.4.34-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.34/jax_rocm60_plugin-0.4.34-cp312-cp312-manylinux_2_28_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jax-v0.4.34.tar.gz

Docker Images

ROCm 6.0.3

docker pull rocm/jax-community:rocm6.0.3-jax0.4.34-py3.10.15
docker pull rocm/jax-community:rocm6.0.3-jax0.4.34-py3.11.10
docker pull rocm/jax-community:rocm6.0.3-jax0.4.34-py3.12.6

ROCm 6.1.3

docker pull rocm/jax-community:rocm6.1.3-jax0.4.34-py3.10.15
docker pull rocm/jax-community:rocm6.1.3-jax0.4.34-py3.11.10
docker pull rocm/jax-community:rocm6.1.3-jax0.4.34-py3.12.6

ROCm 6.2.3

docker pull rocm/jax-community:rocm6.2.3-jax0.4.34-py3.10.15
docker pull rocm/jax-community:latest-py310
docker pull rocm/jax-community:rocm6.2.3-jax0.4.34-py3.11.10
docker pull rocm/jax-community:latest-py311
docker pull rocm/jax-community:rocm6.2.3-jax0.4.34-py3.12.6
docker pull rocm/jax-community:latest-py312
docker pull rocm/jax-community:latest

rocm-jax-v0.4.33

15 Oct 22:32
Compare
Choose a tag to compare

ROCm JAX 0.4.33 Release

ROCm 6.0.3, 6.1.3 and 6.2.3

Python 3.10.15

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jaxlib-0.4.33-cp310-cp310-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jax_rocm60_pjrt-0.4.33-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jax_rocm60_plugin-0.4.33-cp310-cp310-manylinux_2_28_x86_64.whl

Python 3.11.10

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jaxlib-0.4.33-cp311-cp311-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jax_rocm60_pjrt-0.4.33-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jax_rocm60_plugin-0.4.33-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12.6

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jaxlib-0.4.33-cp312-cp312-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jax_rocm60_pjrt-0.4.33-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.33/jax_rocm60_plugin-0.4.33-cp312-cp312-manylinux_2_28_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jax-v0.4.33.tar.gz

Docker Images

ROCm 6.0.3

docker pull rocm/jax-community:rocm6.0.3-jax0.4.33-py3.10.15
docker pull rocm/jax-community:rocm6.0.3-jax0.4.33-py3.11.10
docker pull rocm/jax-community:rocm6.0.3-jax0.4.33-py3.12.6

ROCm 6.1.3

docker pull rocm/jax-community:rocm6.1.3-jax0.4.33-py3.10.15
docker pull rocm/jax-community:rocm6.1.3-jax0.4.33-py3.11.10
docker pull rocm/jax-community:rocm6.1.3-jax0.4.33-py3.12.6

ROCm 6.2.3

docker pull rocm/jax-community:rocm6.2.3-jax0.4.33-py3.10.15
docker pull rocm/jax-community:latest-py310
docker pull rocm/jax-community:rocm6.2.3-jax0.4.33-py3.11.10
docker pull rocm/jax-community:latest-py311
docker pull rocm/jax-community:rocm6.2.3-jax0.4.33-py3.12.6
docker pull rocm/jax-community:latest-py312
docker pull rocm/jax-community:latest

rocm-jaxlib-v0.4.31

14 Oct 14:46
Compare
Choose a tag to compare

ROCm JAX 0.4.31 Release

ROCm 6.0.3, 6.1.3 and 6.2.3

Python 3.10.15

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jaxlib-0.4.31-cp310-cp310-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_pjrt-0.4.31-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_plugin-0.4.31-cp310-cp310-manylinux_2_28_x86_64.whl

Python 3.11.10

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jaxlib-0.4.31-cp311-cp311-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_pjrt-0.4.31-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_plugin-0.4.31-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12.6

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jaxlib-0.4.31-cp312-cp312-manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_pjrt-0.4.31-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_plugin-0.4.31-cp312-cp312-manylinux_2_28_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.31.tar.gz

Docker Images

ROCm 6.0.3

docker pull rocm/jax-community:rocm6.0.3-jax0.4.31-py3.10.15
docker pull rocm/jax-community:rocm6.0.3-jax0.4.31-py3.11.10
docker pull rocm/jax-community:rocm6.0.3-jax0.4.31-py3.12.6

ROCm 6.1.3

docker pull rocm/jax-community:rocm6.1.3-jax0.4.31-py3.10.15
docker pull rocm/jax-community:rocm6.1.3-jax0.4.31-py3.11.10
docker pull rocm/jax-community:rocm6.1.3-jax0.4.31-py3.12.6

ROCm 6.2.3

docker pull rocm/jax-community:rocm6.2.3-jax0.4.31-py3.10.15
docker pull rocm/jax-community:latest-py310
docker pull rocm/jax-community:rocm6.2.3-jax0.4.31-py3.11.10
docker pull rocm/jax-community:latest-py311
docker pull rocm/jax-community:rocm6.2.3-jax0.4.31-py3.12.6
docker pull rocm/jax-community:latest-py312
docker pull rocm/jax-community:latest

rocm-jaxlib-v0.4.30

12 Sep 18:37
Compare
Choose a tag to compare

ROCm jax/jaxlib 0.4.30 Release

ROCm 6.1.1

jaxlib

Python 3.9

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm611-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm611-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm611-cp311-cp311-manylinux2014_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.30.tar.gz

Docker Images

Python 3.9

docker pull rocm/jax-community:rocm6.1.1-jax0.4.30-py3.9.19

Python 3.10

docker pull rocm/jax-community:rocm6.1.1-jax0.4.30-py3.10.14

Python 3.11

docker pull rocm/jax-community:rocm6.1.1-jax0.4.30-py3.11.9

ROCm 6.0.2

jaxlib

Python 3.9

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm602-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm602-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm602-cp311-cp311-manylinux2014_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.30.tar.gz

Docker Images

Python 3.9

docker pull rocm/jax-community:rocm6.0.2-jax0.4.30-py3.9.19

Python 3.10

docker pull rocm/jax-community:rocm6.0.2-jax0.4.30-py3.10.14

Python 3.11

docker pull rocm/jax-community:rocm6.0.2-jax0.4.30-py3.11.9

rocm-jaxlib-v0.4.29

15 Jun 23:54
Compare
Choose a tag to compare

ROCm jax/jaxlib 0.4.29 Release

ROCm 6.1.1

jaxlib

Python 3.9

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.29/jaxlib-0.4.29+rocm611-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.29/jaxlib-0.4.29+rocm611-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.29/jaxlib-0.4.29+rocm611-cp311-cp311-manylinux2014_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.29.tar.gz

Docker Images

Python 3.9

docker pull rocm/jax-community:rocm6.1.1-jax0.4.29-py3.9.19

Python 3.10

docker pull rocm/jax-community:rocm6.1.1-jax0.4.29-py3.10.14

Python 3.11

docker pull rocm/jax-community:rocm6.1.1-jax0.4.29-py3.11.9

ROCm 6.0.2

jaxlib

Python 3.9

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.29/jaxlib-0.4.29+rocm602-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.29/jaxlib-0.4.29+rocm602-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.29/jaxlib-0.4.29+rocm602-cp311-cp311-manylinux2014_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.29.tar.gz

Docker Images

Python 3.9

docker pull rocm/jax-community:rocm6.0.2-jax0.4.29-py3.9.19

Python 3.10

docker pull rocm/jax-community:rocm6.0.2-jax0.4.29-py3.10.14

Python 3.11

docker pull rocm/jax-community:rocm6.0.2-jax0.4.29-py3.11.9

rocm-jaxlib-v0.4.28

23 May 23:53
Compare
Choose a tag to compare

ROCm jax/jaxlib 0.4.28 Release

ROCm 6.1.1

jaxlib

Python 3.9

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.28/jaxlib-0.4.28+rocm611-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.28/jaxlib-0.4.28+rocm611-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.28/jaxlib-0.4.28+rocm611-cp311-cp311-manylinux2014_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.28.tar.gz

Docker Images

Python 3.9

docker pull rocm/jax:rocm6.1.1-jax0.4.28-py3.9.19

Python 3.10

docker pull rocm/jax:rocm6.1.1-jax0.4.28-py3.10.14

Python 3.11

docker pull rocm/jax:rocm6.1.1-jax0.4.28-py3.11.9

ROCm 6.0.2

jaxlib

Python 3.9

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.28/jaxlib-0.4.28+rocm602-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.28/jaxlib-0.4.28+rocm602-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.28/jaxlib-0.4.28+rocm602-cp311-cp311-manylinux2014_x86_64.whl

JAX

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.28.tar.gz

Docker Images

Python 3.9

docker pull rocm/jax:rocm6.0.2-jax0.4.28-py3.9.19

Python 3.10

docker pull rocm/jax:rocm6.0.2-jax0.4.28-py3.10.14

Python 3.11

docker pull rocm/jax:rocm6.0.2-jax0.4.28-py3.11.9

Jax/Jaxlib 0.4.26 release for ROCm

13 Apr 01:36
Compare
Choose a tag to compare

**Jax/Jaxlib 0.4.26 release for ROCm

Note: Added ROCm 6.1 version on May 6th, 2024

ROCm 6.1
Install Jaxlib
Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm610-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm610-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm610-cp311-cp311-manylinux2014_x86_64.whl

Install Jax

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/archive/refs/tags/jaxlib-v0.4.26.tar.gz

Docker Images
Python 3.9

docker pull rocm/jax:rocm6.1.0-jax0.4.26-py3.9.0

Python 3.10

docker pull rocm/jax:rocm6.1.0-jax0.4.26-py3.10.0

Python 3.11

docker pull rocm/jax:rocm6.1.0-jax0.4.26-py3.11.0

ROCm 6.0.0 and 5.7.0
Install Jaxlib
Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm570-cp39-cp39-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm600-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm570-cp310-cp310-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm600-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm570-cp311-cp311-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.26/jaxlib-0.4.26+rocm600-cp311-cp311-manylinux2014_x86_64.whl

Install Jax

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/archive/refs/tags/jaxlib-v0.4.26.tar.gz

Docker Images
Python 3.9

docker pull rocm/jax:rocm5.7.0-jax0.4.26-py3.9.0
docker pull rocm/jax:rocm6.0.0-jax0.4.26-py3.9.0

Python 3.10

docker pull rocm/jax:rocm5.7.0-jax0.4.26-py3.10.0
docker pull rocm/jax:rocm6.0.0-jax0.4.26-py3.10.0

Python 3.11

docker pull rocm/jax:rocm5.7.0-jax0.4.26-py3.11.0
docker pull rocm/jax:rocm6.0.0-jax0.4.26-py3.11.0

Jax/Jaxlib 0.4.24 release for ROCm

27 Feb 23:15
Compare
Choose a tag to compare

**Jax/Jaxlib 0.4.24 release for ROCm

Install Jaxlib
Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm570-cp39-cp39-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm600-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm570-cp310-cp310-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm600-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm570-cp311-cp311-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm600-cp311-cp311-manylinux2014_x86_64.whl

Install Jax

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/archive/refs/tags/jaxlib-v0.4.24.tar.gz

Docker Images
Python 3.9

docker pull rocm/jax:rocm5.7.0-jax0.4.24-py3.9.0
docker pull rocm/jax:rocm6.0.0-jax0.4.24-py3.9.0

Python 3.10

docker pull rocm/jax:rocm5.7.0-jax0.4.24-py3.10.0
docker pull rocm/jax:rocm6.0.0-jax0.4.24-py3.10.0

Python 3.11

docker pull rocm/jax:rocm5.7.0-jax0.4.24-py3.11.0
docker pull rocm/jax:rocm6.0.0-jax0.4.24-py3.11.0

Jax/Jaxlib 0.4.25 release for ROCm

18 Mar 23:08
Compare
Choose a tag to compare

**Jax/Jaxlib 0.4.25 release for ROCm

Install Jaxlib
Python 3.9

python3 -m pip install https://github.com/ROCm/jax/releases/download/jaxlib-v0.4.25/jaxlib-0.4.25+rocm600-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCm/jax/releases/download/jaxlib-v0.4.25/jaxlib-0.4.25+rocm600-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11

python3 -m pip install https://github.com/ROCm/jax/releases/download/jaxlib-v0.4.25/jaxlib-0.4.25+rocm600-cp311-cp311-manylinux2014_x86_64.whl

Install Jax

python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/jaxlib-v0.4.25.tar.gz

Docker Images
Python 3.9

docker pull rocm/jax:rocm6.0.0-jax0.4.25-py3.9.0

Python 3.10

docker pull rocm/jax:rocm6.0.0-jax0.4.25-py3.10.0

Python 3.11

docker pull rocm/jax:rocm6.0.0-jax0.4.25-py3.11.0