feat: Add JAX release support v0.10.0 by erman-gurses · Pull Request #6054 · ROCm/TheRock

erman-gurses · 2026-06-23T07:18:49Z

Summary

ISSUE ID: #3878
Related PR: ROCm/rockrel#65
Extend the JAX multi-arch release workflows to support both rocm-jaxlib-v0.9.1 and rocm-jaxlib-v0.10.0 release lines.

Changes

Release workflow

Add matrix support for:
- rocm-jaxlib-v0.9.1
- rocm-jaxlib-v0.10.0
Add release-line specific metadata:
- jax_repository
- build_mode
- gfx_arch

Build workflow

rocm-jaxlib-v0.9.1

Continue using the existing native build flow.
Build from ROCm/rocm-jax.
Continue using build/ci_build.

rocm-jaxlib-v0.10.0

Build from ROCm/jax.
Build a manylinux image using the release branch Dockerfile from ROCm/rocm-jax.
Build wheels inside the manylinux container using build/build.py.

Testing

On rockrel:

https://github.com/ROCm/rockrel/actions/runs/28217952215/job/83613621212
https://github.com/ROCm/rockrel/actions/runs/28224689859

Preserve existing GPU-family based runner selection through configure_target_run.py.
Add a JAX device visibility sanity check before running tests:

python -c "import jax; print(jax.local_devices())"

- Verified matrix expansion across Python 3.11–3.14.
- Verified native build path for rocm-jaxlib-v0.9.1.
- Verified manylinux build path for rocm-jaxlib-v0.10.0.
- Verified release artifact publication.
- Verified JAX device visibility check executes before tests.

TODO: 
Run JAX builds inside a top-level GitHub Actions job container, similar to existing PyTorch workflows.

therock-pr-bot · 2026-06-23T07:19:45Z

✅ All Policy Checks Passed

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	✅ Pass	—
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

🎉 All policy checks passed!

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

therock-pr-bot · 2026-06-23T07:19:47Z

🎉 All checks passed! This PR is ready for review.

gulsumgudukbay

It would be great if we could move to wheels instead of tarballs and environment variable testing. Unit testing with theRock wheels without any environment variables also allows us to surface packaging issues if there are any. That is how we caught previous packaging issues. This way theRock CI not only tests libraries and systems bugs, it will also test the actual packaging

erman-gurses · 2026-06-25T20:09:14Z

Latest Test: https://github.com/ROCm/rockrel/actions/runs/28196469890 (In Progress)

ScottTodd · 2026-06-25T20:15:22Z

Latest Test: https://github.com/ROCm/rockrel/actions/runs/28196469890 (In Progress)

You can also test dev releases here in TheRock: https://github.com/ROCm/TheRock/actions/workflows/multi_arch_release.yml . We'd prefer that - access to rockrel is intentionally limited as workflow runs there have the potential to push to nightly/prerelease buckets.

ScottTodd

Thanks, looking good now :)

erman-gurses · 2026-06-26T11:03:02Z

Latest Tests are green:
https://github.com/ROCm/rockrel/actions/runs/28217952215/job/83613621212
https://github.com/ROCm/rockrel/actions/runs/28224689859

marbre

I think your changes can go in @erman-gurses but make sure to land this coordinated with your PR to rockrel (whcih still pins your own branch!).

erman-gurses · 2026-06-26T17:20:57Z

I think your changes can go in @erman-gurses but make sure to land this coordinated with your PR to rockrel (whcih still pins your own branch!).

I updated already @marbre, please see the PR: ROCm/rockrel#65

…orkflow (#65) ## Summary Related PR: ROCm/TheRock#6054 Add `rocm_package_find_links_url` to the rockrel JAX release workflow and forward it to the reusable TheRock JAX build workflow. This allows the JAX manylinux build path to consume ROCm packages from a configurable package index instead of relying on a hardcoded URL. ## Changes * Add `rocm_package_find_links_url` workflow input to `multi_arch_release_linux_jax_wheels.yml` * Forward `rocm_package_find_links_url` to the reusable `multi_arch_build_linux_jax_wheels.yml` workflow * Default the value to the current dev multi-arch package index: * `https://rocm.devreleases.amd.com/whl-multi-arch/` ## Motivation The JAX manylinux build flow now installs ROCm packages from a package index. Passing the package source as an input allows the workflow to support different release channels (dev, nightly, prerelease) without hardcoding a specific index URL in the build workflow. Tests: https://github.com/ROCm/rockrel/actions/runs/28086158341

## Summary ISSUE ID: #3878 Related PR: ROCm/rockrel#65 Extend the JAX multi-arch release workflows to support both `rocm-jaxlib-v0.9.1` and `rocm-jaxlib-v0.10.0` release lines. ## Changes ### Release workflow - Add matrix support for: - `rocm-jaxlib-v0.9.1` - `rocm-jaxlib-v0.10.0` - Add release-line specific metadata: - `jax_repository` - `build_mode` - `gfx_arch` ### Build workflow #### rocm-jaxlib-v0.9.1 - Continue using the existing native build flow. - Build from `ROCm/rocm-jax`. - Continue using `build/ci_build`. #### rocm-jaxlib-v0.10.0 - Build from `ROCm/jax`. - Build a manylinux image using the release branch Dockerfile from `ROCm/rocm-jax`. - Build wheels inside the manylinux container using `build/build.py`. ### Testing On rockrel: https://github.com/ROCm/rockrel/actions/runs/28217952215/job/83613621212 https://github.com/ROCm/rockrel/actions/runs/28224689859 - Preserve existing GPU-family based runner selection through `configure_target_run.py`. - Add a JAX device visibility sanity check before running tests: ```bash python -c "import jax; print(jax.local_devices())" - Verified matrix expansion across Python 3.11–3.14. - Verified native build path for rocm-jaxlib-v0.9.1. - Verified manylinux build path for rocm-jaxlib-v0.10.0. - Verified release artifact publication. - Verified JAX device visibility check executes before tests. TODO: Run JAX builds inside a top-level GitHub Actions job container, similar to existing PyTorch workflows.

## Summary Incremental follow-up to #6054 (JAX 0.10.0 support). Adds the JAX 0.10.2 release to the multi-arch Linux JAX wheel build/test matrix. - Adds `rocm-jaxlib-v0.10.2` to the release matrix, wired identically to 0.10.0 (`ROCm/jax`, `manylinux` build mode, `device-all`). - Uses the **default** XLA allocator (`XLA_PYTHON_CLIENT_ALLOCATOR=default`) for the 0.10.2 test cell only; 0.9.1 and 0.10.0 keep `platform`. - Documents 0.10.2 in the JAX support table. ## Test plan - [ ] `Multi-Arch Release Linux JAX Wheels` 0.10.2 cell builds wheels - [ ] JAX test subset passes with the default allocator I ran full matrix run got all green : https://github.com/ROCm/TheRock/actions/runs/28336527815/job/83947970257

Cherry-picks 3b10211 into the 7.14 release branch. --- ## Summary ISSUE ID: #3878 Related PR: ROCm/rockrel#65 Extend the JAX multi-arch release workflows to support both `rocm-jaxlib-v0.9.1` and `rocm-jaxlib-v0.10.0` release lines. ## Changes ### Release workflow - Add matrix support for: - `rocm-jaxlib-v0.9.1` - `rocm-jaxlib-v0.10.0` - Add release-line specific metadata: - `jax_repository` - `build_mode` - `gfx_arch` ### Build workflow #### rocm-jaxlib-v0.9.1 - Continue using the existing native build flow. - Build from `ROCm/rocm-jax`. - Continue using `build/ci_build`. #### rocm-jaxlib-v0.10.0 - Build from `ROCm/jax`. - Build a manylinux image using the release branch Dockerfile from `ROCm/rocm-jax`. - Build wheels inside the manylinux container using `build/build.py`. ### Testing On rockrel: https://github.com/ROCm/rockrel/actions/runs/28217952215/job/83613621212 https://github.com/ROCm/rockrel/actions/runs/28224689859 - Preserve existing GPU-family based runner selection through `configure_target_run.py`. - Add a JAX device visibility sanity check before running tests: ```bash python -c "import jax; print(jax.local_devices())" - Verified matrix expansion across Python 3.11–3.14. - Verified native build path for rocm-jaxlib-v0.9.1. - Verified manylinux build path for rocm-jaxlib-v0.10.0. - Verified release artifact publication. - Verified JAX device visibility check executes before tests. TODO: Run JAX builds inside a top-level GitHub Actions job container, similar to existing PyTorch workflows. Co-authored-by: erman-gurses <99776114+erman-gurses@users.noreply.github.com>

…orkflow (#65) (#68) Cherry-pick b775c04 into the 7.14 release branch. --- ## Summary Related PR: ROCm/TheRock#6054 Add `rocm_package_find_links_url` to the rockrel JAX release workflow and forward it to the reusable TheRock JAX build workflow. This allows the JAX manylinux build path to consume ROCm packages from a configurable package index instead of relying on a hardcoded URL. ## Changes * Add `rocm_package_find_links_url` workflow input to `multi_arch_release_linux_jax_wheels.yml` * Forward `rocm_package_find_links_url` to the reusable `multi_arch_build_linux_jax_wheels.yml` workflow * Default the value to the current dev multi-arch package index: * `https://rocm.devreleases.amd.com/whl-multi-arch/` ## Motivation The JAX manylinux build flow now installs ROCm packages from a package index. Passing the package source as an input allows the workflow to support different release channels (dev, nightly, prerelease) without hardcoding a specific index URL in the build workflow. Tests: https://github.com/ROCm/rockrel/actions/runs/28086158341 Co-authored-by: erman-gurses <99776114+erman-gurses@users.noreply.github.com>

## Summary Issue ID: #6226 Updated the JAX section in `RELEASES.md` to reflect current supported versions and installation methods. ### Changes - **Multi-arch JAX section**: Added clear installation instructions using the unified index (`https://rocm.nightlies.amd.com/whl-multi-arch/`) with `[device-*]` extras. - Updated supported JAX versions to **0.9.1** and **0.10.0**. - Added note about version pinning (`jax`, `jax_rocm7_plugin`, `jax_rocm7_pjrt`). - Improved consistency with the PyTorch documentation section. - Minor cleanup and better cross-references. ### Related - #6054 (JAX 0.10.0 support) - https://github.com/ROCm/rockrel/actions/runs/28346375566 (successful nightly) This makes the documentation accurate and user-friendly for both 0.9.1 and 0.10.0 users.

Add JAX release support v0.10.0

a4ca6a2

github-project-automation Bot added this to TheRock Triage Jun 23, 2026

github-project-automation Bot moved this to TODO in TheRock Triage Jun 23, 2026

therock-pr-bot Bot added the Not ready to Review PR has unresolved policy failures — reviews blocked label Jun 23, 2026

erman-gurses requested review from ScottTodd and marbre June 23, 2026 07:19

erman-gurses added the ci:skip Skip all CI builds/tests for this PR label Jun 23, 2026

erman-gurses linked an issue Jun 23, 2026 that may be closed by this pull request

[CI] Build and test JAX Python packages as part of ci.yml #3878

Closed

erman-gurses changed the title ~~Add JAX release support v0.10.0~~ [Feature] Add JAX release support v0.10.0 Jun 23, 2026

erman-gurses changed the title ~~[Feature] Add JAX release support v0.10.0~~ feat - Add JAX release support v0.10.0 Jun 23, 2026

Copilot started work on behalf of erman-gurses June 23, 2026 07:29 View session

erman-gurses changed the title ~~feat - Add JAX release support v0.10.0~~ feat: Add JAX release support v0.10.0 Jun 23, 2026

therock-pr-bot Bot removed the Not ready to Review PR has unresolved policy failures — reviews blocked label Jun 23, 2026

Copilot finished work on behalf of erman-gurses June 23, 2026 07:38

Fix gfx_arch: all

dd528cb

erman-gurses requested review from gulsumgudukbay and tsrw2048 and removed request for marbre June 23, 2026 16:16

ScottTodd requested a review from marbre June 23, 2026 17:40

ScottTodd reviewed Jun 23, 2026

View reviewed changes

gulsumgudukbay requested changes Jun 23, 2026

View reviewed changes

ScottTodd mentioned this pull request Jun 23, 2026

Modify TheRock manylinux Dockerfile to take an arch build arg. ROCm/rocm-jax#470

Merged

erman-gurses added 6 commits June 23, 2026 20:31

Remove tarball

baef7fe

Test

6fc71c7

Test

74ab4d7

Add back tarball for 0.9.0

19eda7c

Test

c170950

Test

e6f750d

tsrw2048 approved these changes Jun 25, 2026

View reviewed changes

Black

bf06e82

ScottTodd approved these changes Jun 25, 2026

View reviewed changes

Comment thread .github/workflows/multi_arch_build_linux_jax_wheels.yml Outdated

erman-gurses added 2 commits June 25, 2026 20:52

Remove the debug code

ca2eaed

Use --find-links for ROCm meta-package installation in manylinux build

9ce38fd

therock-pr-bot Bot added the Not ready to Review PR has unresolved policy failures — reviews blocked label Jun 26, 2026

Remove added package_index_url

6b9b464

therock-pr-bot Bot removed the Not ready to Review PR has unresolved policy failures — reviews blocked label Jun 26, 2026

erman-gurses and others added 3 commits June 26, 2026 03:30

Workaround for --find-links

139df3c

Merge branch 'main' into users/erman-gurses/add-jax10x-support

31e9253

Revert workaround code and use ROCM_PACKAGE_FIND_LINKS_URL

da73551

marbre reviewed Jun 26, 2026

View reviewed changes

erman-gurses requested review from marbre June 26, 2026 17:22

erman-gurses merged commit 3b10211 into main Jun 26, 2026
136 of 144 checks passed

erman-gurses deleted the users/erman-gurses/add-jax10x-support branch June 26, 2026 17:44

github-project-automation Bot moved this from TODO to Done in TheRock Triage Jun 26, 2026

magaonka-amd mentioned this pull request Jun 28, 2026

feat: Add JAX release support v0.10.2 #6195

Merged

2 tasks

marbre mentioned this pull request Jun 29, 2026

Pass ROCm package rocm_package_find_links_url through JAX release workflow (#65) ROCm/rockrel#68

Merged

erman-gurses mentioned this pull request Jun 29, 2026

docs: update JAX installation instructions in RELEASES.md #6225

Merged

Uh oh!

Conversation

erman-gurses commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Release workflow

Build workflow

rocm-jaxlib-v0.9.1

rocm-jaxlib-v0.10.0

Testing

Uh oh!

therock-pr-bot Bot commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ All Policy Checks Passed

Uh oh!

therock-pr-bot Bot commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gulsumgudukbay left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

erman-gurses commented Jun 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ScottTodd commented Jun 25, 2026

Uh oh!

ScottTodd left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

erman-gurses commented Jun 26, 2026

Uh oh!

marbre left a comment

Choose a reason for hiding this comment

Uh oh!

erman-gurses commented Jun 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

erman-gurses commented Jun 23, 2026 •

edited

Loading

therock-pr-bot Bot commented Jun 23, 2026 •

edited

Loading

therock-pr-bot Bot commented Jun 23, 2026 •

edited

Loading

gulsumgudukbay left a comment •

edited

Loading

erman-gurses commented Jun 25, 2026 •

edited

Loading