Jax/Jaxlib 0.4.13 release for ROCm 5.5

Jax/Jaxlib 0.4.13 release for ROCm 5.5.0

Install Jaxlib
Python 3.8

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.13-rocm550/jaxlib-0.4.13+rocm550-cp38-cp38-manylinux2014_x86_64.whl

Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.13-rocm550/jaxlib-0.4.13+rocm550-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.13-rocm550/jaxlib-0.4.13+rocm550-cp310-cp310-manylinux2014_x86_64.whl

Install Jax

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

Docker Images
Python 3.8

docker pull rocm/jax:rocm5.5.0-jax0.4.13-py3.8.0

Python 3.9

docker pull rocm/jax:rocm5.5.0-jax0.4.13-py3.9.0

Python 3.10

docker pull rocm/jax:rocm5.5.0-jax0.4.13-py3.10.0

Jax/Jaxlib 0.4.12 release for ROCm

Jax/Jaxlib 0.4.12 release for ROCm 5.5.0

Install Jaxlib
Python 3.8

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.12-rocm550/jaxlib-0.4.12+rocm550-cp38-cp38-manylinux2014_x86_64.whl

Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.12-rocm550/jaxlib-0.4.12+rocm550-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.12-rocm550/jaxlib-0.4.12+rocm550-cp310-cp310-manylinux2014_x86_64.whl

Install Jax

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

Docker Images
Python 3.8

docker pull rocm/jax:rocm5.5.0-jax0.4.12-py3.8.0

Python 3.9

docker pull rocm/jax:rocm5.5.0-jax0.4.12-py3.9.0

Python 3.10

docker pull rocm/jax:rocm5.5.0-jax0.4.12-py3.10.0

Jax/Jaxlib 0.4.11 release for ROCm

Jax/Jaxlib 0.4.11 release for ROCm 5.5.0

Install Jaxlib
Python 3.8

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.11-rocm550/jaxlib-0.4.11+rocm550-cp38-cp38-manylinux2014_x86_64.whl

Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.11-rocm550/jaxlib-0.4.11+rocm550-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.11-rocm550/jaxlib-0.4.11+rocm550-cp310-cp310-manylinux2014_x86_64.whl

Install Jax

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

Docker Images
Python 3.8

docker pull rocm/jax:rocm5.5.0-jax0.4.11-py3.8.0

Python 3.9

docker pull rocm/jax:rocm5.5.0-jax0.4.11-py3.9.0

Python 3.10

docker pull rocm/jax:rocm5.5.0-jax0.4.11-py3.10.0

Jax/Jaxlib 0.4.10 release for ROCm

Jax/Jaxlib 0.4.10 release for ROCm 5.5.0

Install Jaxlib
Select Python version (3.8, 3.9 or 3.10)

Python 3.8

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-0.4.10-rocm550/jaxlib-0.4.10+rocm550-cp38-cp38-manylinux2014_x86_64.whl

Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-0.4.10-rocm550/jaxlib-0.4.10+rocm550-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-0.4.10-rocm550/jaxlib-0.4.10+rocm550-cp310-cp310-manylinux2014_x86_64.whl

Install Jax
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/archive/refs/tags/jaxlib-0.4.10-rocm550.tar.gz

Docker Images
Python 3.8

docker pull rocm/jax:rocm5.5.0-jax0.4.10-py3.8.0

Python 3.9

docker pull rocm/jax:rocm5.5.0-jax0.4.10-py3.9.0

Python 3.10

docker pull rocm/jax:rocm5.5.0-jax0.4.10-py3.10.0

Jax/Jaxlib 0.4.9 release for ROCm 5.5.0

Jax/Jaxlib 0.4.9 release for ROCm

Install Jaxlib
Python 3.8

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-0.4.9-rocm550/jaxlib-0.4.9.550-cp38-cp38-manylinux2014_x86_64.whl

Python 3.9

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-0.4.9-rocm550/jaxlib-0.4.9.550-cp39-cp39-manylinux2014_x86_64.whl

Python 3.10

python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-0.4.9-rocm550/jaxlib-0.4.9.550-cp310-cp310-manylinux2014_x86_64.whl

Install Jax

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

Docker Images
Python 3.8

docker pull rocm/jax:rocm5.5.0-jax0.4.9.550-py3.8.0

Python 3.9

docker pull rocm/jax:rocm5.5.0-jax0.4.9.550-py3.9.0

Python 3.10

docker pull rocm/jax:rocm5.5.0-jax0.4.9.550-py3.10.0

jaxlib-v0.4.6-rocm550

Update docker file to ROCm 5.5