What's Changed

• Llama Performance Test by @psanal35 in #141
• The ROCm plugin should fail gracefully by @JehandadKhan in #143
• Update workflows (ci, and nightly) for ROCm 7.0.2 by @psanal35 in #150
• Update Llama performance workflow by adding a new runner, a model, and JAX version support by @psanal35 in #149
• Pull image before tests in nightly-CI so docker inspect always works by @psanal35 in #156
• Fix build number for DLM performance CI and update for JAX 0.7.1 by @psanal35 in #158
• Add /opt/rocm/lib to RPATH so missing ROCm libs are found by @psanal35 in #147
• Explicitly install missing libdw1 library for ROCm 7.0.x before running tests in CI by @psanal35 in #160
• Add explicit permissions for nightly by @charleshofer in #155
• Use wheels from latest scheduled nightly run in Llama perf by @psanal35 in #164
• Sync skip.ini usage on master with updates from v0.7.0 release by @psanal35 in #167
• Move to JAX 0.8.0 and ROCm 7.2 by @charleshofer in #173
• Port change for jaxlib wheel build to 0.8.0 by @charleshofer in #183
• JAX Wheel Tagging with version extra by @subodh-dubey-amd in #184
• Bump dev setup build number for ROCm 7.2.0 for 0.8.0 by @charleshofer in #186
• Port run CI on release branches for 0.8.0 by @charleshofer in #192
• Add ROCm version suffix to jaxlib by @charleshofer in #191
• Bump XLA commit for 0.8.0 XLA ports by @charleshofer in #202
• Update JAX version used in nightly workflow for the release branch by @psanal35 in #204
• Bump XLA commit for 0.8.0 by @charleshofer in #211
• Use ROCm RC14 on 0.8.0 by @charleshofer in #213
• Bump commit for profiling changes on 0.8.0 by @charleshofer in #216
• Bump XLA commit for fusion emitter and silu fixes by @charleshofer in #220
• Bump XLA commit for hipSolver fix by @charleshofer in #224
• Remove buggy device count patch by @charleshofer in #227
• Bump XLA for collectives fix by @charleshofer in #229
• Update ROCm 7.2.0 build number by @charleshofer in #230
• Add all python versions to nightly by @charleshofer in #231
• Fix py version typo by @charleshofer in #232

Full Changelog: rocm-jax-v0.7.1...rocm-jax-v0.8.0

JAX 0.6.2 Installation Instructions

Python 3.10

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jaxlib-0.6.2-cp310-cp310-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_pjrt-0.6.2-py3-none-manylinux_2_28_x86_64.whl  https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_plugin-0.6.2-cp310-cp310-manylinux_2_28_x86_64.whl

Python 3.11

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jaxlib-0.6.2-cp311-cp311-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_pjrt-0.6.2-py3-none-manylinux_2_28_x86_64.whl  https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_plugin-0.6.2-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jaxlib-0.6.2-cp312-cp312-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_pjrt-0.6.2-py3-none-manylinux_2_28_x86_64.whl  https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_plugin-0.6.2-cp312-cp312-manylinux_2_28_x86_64.whl

Python 3.13

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jaxlib-0.6.2-cp313-cp313-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_pjrt-0.6.2-py3-none-manylinux_2_28_x86_64.whl  https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.2/jax_rocm7_plugin-0.6.2-cp313-cp313-manylinux_2_28_x86_64.whl

What's Changed

• Fix devenv build on 0.6.0 by @mminutoli in #144
• Remove the remaining namespace patches as they are no longer needed by @psanal35 in #146
• cherry-pick 73af0ea to rocm-jaxlib-v0.6.0 by @jujaykka in #131

New Contributors

• @mminutoli made their first contribution in #144

Full Changelog: rocm-jax-v0.6.0...rocm-jax-v0.6.2

JAX 0.7.1 Installation Instructions

Python 3.11

jaxlib

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

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl  https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_plugin-0.7.1-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12

jaxlib

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

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl   https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_plugin-0.7.1-cp312-cp312-manylinux_2_28_x86_64.whl

Python 3.13

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jaxlib-0.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jaxlib-0.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl  https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_plugin-0.7.1-cp313-cp313-manylinux_2_28_x86_64.whl

What's Changed

• CI: push docker image with perf-maxtext tag by @Ruturaj4 in #44
• Add CI status badge to README by @charleshofer in #62
• Move maxtext dockerfile to ci directory by @charleshofer in #65
• Bugfix a bazel debug rule and update gitignore (#64) by @mrodden in #66
• Move to the latest master at 44f7d8796da5b42f3b9483fd6750edaf8018f538 by @Ruturaj4 in #67
• added three args to docker dev setup script by @gulsumgudukbay in #69
• Bump Bazel to 7.4.1 to match JAX proper & XLA by @Arech8 in #70
• Generate Makefile with dynamic AMDGPU_TARGETS using rocminfo by @Ruturaj4 in #76
• Blackify and Pylintify get_rocm.py by @charleshofer in #60
• Add gfx950 to build.py by @Ruturaj4 in #78
• Update Ub24 dockerfile for externally managed python issue by @pramenku in #77
• adding update to get_rocm.py script to update the apt packages by @gulsumgudukbay in #79
• Enable builds with TheRock by @charleshofer in #58
• Clean up week-old docker things by @charleshofer in #86
• Get the nightly workflow to pass its unit tests by @charleshofer in #63
• Tell the docker build script to tag builds as 'nightly' by @charleshofer in #87
• Bump XLA version by @Ruturaj4 in #90
• Adding abort support to UT suite and migrating multiple-GPU requiring tests. by @gulsumgudukbay in #91
• Fixed test runners by @gulsumgudukbay in #94
• [master] fix wheel module not found error by @kiran-thumma in #99
• Add documentation on how to build rocm-jax and set up the dev env by @charleshofer in #101
• Add pip install to Ubuntu 22 docker build by @charleshofer in #102
• Improve stack.py develop options by @charleshofer in #98
• Use test scripts from the rocm/rocm-jax repo rather than rocm/jax by @charleshofer in #96
• Upload pytest JSON results to MySQL for nightly runs by @psanal35 in #75
• Fix build number for ROCm 7.0.0 by @psanal35 in #108
• Bugfix stack.py to re-enable ./xla checkout, +improvements by @Arech8 in #112
• Add --continue_on_fail flag to run_multi_gpu.py by @psanal35 in #114
• Add upload wheels by @Ruturaj4 in #116
• Prep for Jax v0.7.1 release by @zahiqbal in #105
• fix to work with released ROCm 7.0.0 when using --rocm-version=7.0.0 by @jujaykka in #117
• CI: Prevent intermittent checkout failures on NFS-based runners by @leo-amd in #119
• Fix nightly and use ROCm 7.0.0 for PRs by @charleshofer in #122
• Detect ROCm version number for the wheels from the actual ROCm version number by @gabeweisz in #121
• Used quiet^2 mode for pytest collect process by @gulsumgudukbay in #126
• Remove verbose logging from auditwheel by @charleshofer in #100
• Add test skips file by @charleshofer in #97
• Allow user-provided ROCm and calculate namespace correctly in more places by @gabeweisz in #127
• Single-GPU UT script path fix by @gulsumgudukbay in #133
• Handle tests/ prefix in pytest nodeid parsing by @psanal35 in #134
• Updated XLA pinned version by @zahiqbal in #138
• Init Llama performance test by @psanal35 in #140
• Bake commit info wheels by @AratiGanesh in #104
• Added jaxlib wheel build infrastructure and included it in docker-images by @zahiqbal in #142
• updated xla commit on repo rocm/xla - rocm-jaxlib-v0.7.1 branch by @zahiqbal in #157
• Adding rpath for /opt/rocm/libs since librccl.so etc is not found (#135) by @zahiqbal in #159
• Fixed multi-gpu tests deselection by @zahiqbal in #161
• Pull image before tests in nightly-CI so docker inspect always works by @zahiqbal in #162
• Updated XLA pinned to tip of rocm/xla, branch rocm-jaxlib-v0.7.1 by @zahiqbal in #166

New Contributors

• @kiran-thumma made their first contribution in #99
• @zahiqbal made their first contribution in #105
• @leo-amd made their first contribution in #119
• @gabeweisz made their first contribution in #121
• @AratiGanesh made their first contribution in #104

Full Changelog: rocm-jax-v0.6.0...rocm-jax-v0.7.1

This is the first release of the JAX Plugin Repository. This repository contains sources for the ROCm plugin for JAX, as well as Dockerfiles used to build AMD's rocm/jax images. The main difference is that this repository generates two wheels (plugin and PJRT wheels for ROCm JAX). These wheels can be found as artifacts of the 0.6.0 release.

Installation Instructions

Python 3.10.12

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp310-cp310-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp310-cp310-manylinux_2_28_x86_64.whl

Python 3.11.0

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp311-cp311-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12.3

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp312-cp312-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp312-cp312-manylinux_2_28_x86_64.whl

Python 3.13.8

jaxlib

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.0-cp313-cp313-manylinux2014_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_pjrt-0.6.0-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jax_rocm7_plugin-0.6.0-cp313-cp313-manylinux_2_28_x86_64.whl

jaxlib 0.6.2 Installation Instructions

Python 3.10.12

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp310-cp310-manylinux2014_x86_64.whl

Python 3.11.0

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp311-cp311-manylinux2014_x86_64.whl

Python 3.12.3

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp312-cp312-manylinux2014_x86_64.whl

Python 3.13.8

python3 -m pip install https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.6.0/jaxlib-0.6.2-cp313-cp313-manylinux2014_x86_64.whl

What's Changed

• Add upstream CI watcher workflow and script by @charleshofer in #3
• Update build scripts for new plugin layout by @mrodden in #1
• Add setup instructions and scripts by @mrodden in #5
• GH-90: Add workflow files for test CI jobs by @charleshofer in #4
• Add blacken script and format fixes by @mrodden in #9
• Add linters for Python and Bash by @charleshofer in #2
• Add .yamllint.yml config file by @charleshofer in #12
• Change clang install to use packages by @mrodden in #11
• Fix upstream CI watcher job by @charleshofer in #14
• Improve Black output with diff and color options by @psanal35 in #17
• Add .shellcheckrc file by @charleshofer in #13
• Fix bugs in workflow YML files by @charleshofer in #8
• Add changes for compatibility with jax 0.6.0 by @mrodden in #18
• patch for using HIPBlas_V2 by @gulsumgudukbay in #21
• fixed the typo for jax ref by @gulsumgudukbay in #20
• Add build fixes for version 0.6.0 by @mrodden in #22
• Add filter parameter for build_dockers target by @mrodden in #27
• update stack.py to reflect ROCm 7.0.0 by @gulsumgudukbay in #26
• update ci yml file for jax 0.6.0 by @gulsumgudukbay in #28
• Use internal ROCm builds for nightly test runs by @charleshofer in #23
• Add scripts to run linters locally by @charleshofer in #25
• Add group setting to nightly CI concurrency by @charleshofer in #32
• Update function to accurately create ROCm GPU targets and touch versi… by @Ruturaj4 in #36
• Add step to make sure we can clean up old files by @charleshofer in #38
• Explicitly pass the mi-250 label for nightly runs by @charleshofer in #33
• Add namespace parameter to makefile template by @mrodden in #31
• Add performance script by @Ruturaj4 in #24
• Clean up the cleanup container by @charleshofer in #40
• Restrict perf job to run nightly, or manually by @Ruturaj4 in #42
• Bump ROCm 7.0 build number by @charleshofer in #45
• format build.py by @Ruturaj4 in #46
• Remove continue on error and fix invalid workflow by @charleshofer in #49
• Add gfx950 arch by @Ruturaj4 in #37
• Add option to remove docker by @charleshofer in #43
• Adding namespace patch for correctly finding kernels if ROCm 7 is used by @gulsumgudukbay in #50
• add rocm-smi and rocminfo to check gpu info by @Ruturaj4 in #47
• Fix Docker patch step for ROCm >= 7.0: escape by @Ruturaj4 in #51
• Remove process removal using kill by @Ruturaj4 in #53
• Namespace patch changes for setup instructions and docker_dev_setup script by @gulsumgudukbay in #52
• Add nightly docker tag and pull instructions by @charleshofer in #54
• adding CI script UT instructions to the README by @gulsumgudukbay in #56
• Add script to upload MaxText summary metrics to MySQL ci_model_runs table by @Ruturaj4 in #57
• Bugfix a bazel debug rule and update gitignore by @Arech8 in #64
• Update XLA submodule to latest commit on rocm/xla:rocm-jaxlib-v0.6.0 (44f7d8796da5) by @Ruturaj4 in #68
• Add gfx950 to build.py (#78) by @Ruturaj4 in #80
• [rocm-jaxlib-v0.6.0] CP PR#99 fix wheel module not found error by @pramenku in #103
• Update ROCm build info in nightly.yml by @gulsumgudukbay in #106
• Cherry-picked from #96 only nightly modifications to fix permission issue by @gulsumgudukbay in #109
• Update XLA commit hash to the recent commit on 9/25/2025 by @gulsumgudukbay in #128
• Adding rpath for /opt/rocm/libs since librccl.so etc is not found by @gulsumgudukbay in #135
• adding last polishes to wheels by @gulsumgudukbay in #136

Full Changelog: https://github.com/ROCm/rocm-jax/commits/rocm-jax-v0.6.0