Releases: ROCm/jax
Releases · ROCm/jax
rocm-jax-v0.6.0
What's Changed
- Add workflow for nightly pull from upstream by @charleshofer in #115
- Allow devs to kick off sync job manually by @charleshofer in #119
- Unpin container in CI build and remove libssl-dev install by @charleshofer in #123
- Rename the CI flow to 'ROCm CI' and only run it on PRs to rocm-main by @charleshofer in #126
- Fix nightly sync permissions by @charleshofer in #124
- Add GHA workflow for opening PRs upstream by @charleshofer in #116
- Use middle branch for upstream sync by @charleshofer in #128
- Fix upstream sync checkout by @charleshofer in #130
- Add git fetch by @charleshofer in #132
- Make daily sync permissions at the workflow level and fix merge CI by @charleshofer in #143
- Fix upstream PR workflow to use origin branches by @charleshofer in #151
- Add token to open upstream PRs by @charleshofer in #152
- Change the workflow for opening upstream PRs to post links that open PRs by @charleshofer in #157
- Fix rebase command to exclude rocm-main by @charleshofer in #158
- Fix user identity for rebase by @charleshofer in #159
- Fix the link to the downstream PR by @charleshofer in #160
- Use the reference format for links instead of inline by @charleshofer in #162
- jax-mlGH-9948: Fix kernel build by @charleshofer in #173
- GH 9948: Automerge daily sync PRs by @charleshofer in #181
- Add upload wheels file for pypi by @Ruturaj4 in #184
- Add GPU CI by @charleshofer in #137
- Run less tests for CI by @charleshofer in #199
- Add dev guide by @charleshofer in #188
- Run GPU CI on QA PRs by @charleshofer in #228
- Use a GitHub app for syncing rocm-main and upstream main by @charleshofer in #224
- Add CODEOWNERS file by @charleshofer in #236
- Use bazel for PR tests by @charleshofer in #216
- Change CODEOWNERS by @charleshofer in #237
- removing pytest-csv from pytests by @zahiqbal in #248
- Add gfx1101 Target to Jax builds by @JehandadKhan in #249
- Use XLA with include fix by @charleshofer in #256
- Fix C++23 build errors by @charleshofer in #257
- Deprecate obsolete gfx versions by @JehandadKhan in #273
- Add JSON output to multi-GPU tests by @charleshofer in #274
- Count test totals correctly for dashboards by @charleshofer in #280
- Fix auditwheel version issue by @mrodden in #288
- Fix ROCm build README by @charleshofer in #284
- Enable upstream CI on release by @charleshofer in #296
- Add back raw totals in JSON reports by @charleshofer in #281
- Fixes for 0.5.0 build ported to rocm-main by @mrodden in #315
- Use rebase for upstream sync by @charleshofer in #325
- Add Python and ROCm version matrix for rocm-main by @charleshofer in #314
- Port Jenkins release CI to Actions by @charleshofer in #336
- Fix backslash and quoting errors in release CI by @charleshofer in #374
- Remove option for CI to use custom XLA by @charleshofer in #365
- adding jax unit test tridiagonal for rocm platform by @zahiqbal in #369
- [0.6.0-UT] Cherry pick UTs from 0.5.0 by @gulsumgudukbay in #438
- [0.6.0-UT] Unskip ut 0.6.0 by @gulsumgudukbay in #439
- [0.6.0-UT] Add shared utility function get_rocm_version to test_util.py by @psanal35 in #445
- [0.6.0-UT] Skip test_dot for ROCm <6.5, and shapes exceed shared memory limits by @psanal35 in #448
- [0.6.0-UT] adding a QA test requirements file for the run_single_gpu test script… by @gulsumgudukbay in #446
- [0.6.0-UT] Fix rocm-test-requirements by @gulsumgudukbay in #452
- [0.6.0-UT] Select the platform GPUs for a few MultiProcessGpuTest test cases. by @phambinhfin in #443
- [0.6.0-UT] Fix NumpyLinalgTest test cases by @phambinhfin in #453
- Build with ROCm 7.0 on new runner by @charleshofer in #449
- updated the XLA commit hash by @gulsumgudukbay in #455
- Fix get_rocm_version to skip tests if installation is missing by @psanal35 in #456
- Add missing header
warnings
for test_util.py by @psanal35 in #458 - Remove warnings.warn from get_rocm_version by @psanal35 in #459
- [0.6.0-UT] skip failing tests by @gulsumgudukbay in #442
Full Changelog: rocm-jax-v0.5.0...rocm-jax-v0.6.0
rocm-jax-v0.5.0
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
rocm-jax-v0.4.35
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
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
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
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
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
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
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
**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