Skip to content

Jax/Jaxlib 0.4.26 release for ROCm

Compare
Choose a tag to compare
@Ruturaj4 Ruturaj4 released this 13 Apr 01:36

**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