Skip to content

rocm-jax-v0.5.0

Latest
Latest
Compare
Choose a tag to compare
@Ruturaj4 Ruturaj4 released this 24 Jan 22:33
· 3 commits to rocm-jaxlib-v0.5.0 since this release
rocm-jax-v0.5.0
cdf1ef5
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

ROCm JAX 0.5.0 Release

ROCm 6.0.3, 6.2.4 and 6.3.1

Python 3.10.16

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.5.0/jaxlib-0.5.0-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.5.0/jax_rocm60_pjrt-0.5.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.5.0/jax_rocm60_plugin-0.5.0-cp310-cp310-manylinux_2_28_x86_64.whl

Python 3.11.11

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.5.0/jaxlib-0.5.0-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.5.0/jax_rocm60_pjrt-0.5.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.5.0/jax_rocm60_plugin-0.5.0-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12.8

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.5.0/jaxlib-0.5.0-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.5.0/jax_rocm60_pjrt-0.5.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.5.0/jax_rocm60_plugin-0.5.0-cp312-cp312-manylinux_2_28_x86_64.whl

JAX

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

Docker Images

ROCm 6.0.3

docker pull rocm/jax-community:rocm6.0.3-jax0.5.0-py3.10.16
docker pull rocm/jax-community:rocm6.0.3-jax0.5.0-py3.11.11
docker pull rocm/jax-community:rocm6.0.3-jax0.5.0-py3.12.8

ROCm 6.2.4

docker pull rocm/jax-community:rocm6.2.4-jax0.5.0-py3.10.16
docker pull rocm/jax-community:rocm6.2.4-jax0.5.0-py3.11.11
docker pull rocm/jax-community:rocm6.2.4-jax0.5.0-py3.12.8

ROCm 6.3.1

docker pull rocm/jax-community:rocm6.3.1-jax0.5.0-py3.10.16
docker pull rocm/jax-community:latest-py310
docker pull rocm/jax-community:rocm6.3.1-jax0.5.0-py3.11.11
docker pull rocm/jax-community:latest-py311
docker pull rocm/jax-community:rocm6.3.1-jax0.5.0-py3.12.8
docker pull rocm/jax-community:latest
docker pull rocm/jax-community:latest-py312
Assets 9
Loading