Add optional JAX release build toggle by ScottTodd · Pull Request #6117 · ROCm/TheRock

ScottTodd · 2026-06-24T21:48:08Z

Motivation

Progress on #5634 and #6218.

Follow-up to these PRs:

This brings PyTorch and JAX back into alignment with optional triggering, a script to define the job matrix, and the ability to trigger builds for a single python version at a time to save CI resources.

Technical Details

Important

The python_version input will need to be added to https://github.com/ROCm/rockrel/blob/main/.github/workflows/multi_arch_release_linux_jax_wheels.yml, like with ROCm/rockrel#54 and ROCm/rockrel#55.

I've prepared this PR for that: ROCm/rockrel#70.

Test Plan

New unit tests
Trigger a dev release with multi_arch_release_linux.yml, setting build_jax to false
- https://github.com/ROCm/TheRock/actions/runs/28469583076
Trigger multi_arch_release_linux_jax_wheels.yml for a single python_version
- https://github.com/ROCm/TheRock/actions/runs/28469829633

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

Mirror the PyTorch release workflow toggle for JAX so manual multi-arch dev releases can skip JAX wheels by default while workflow_call release runs continue to build them. Add a small JAX release matrix script so JAX can share the top-level python_version narrowing behavior. Also add focused tests for the JAX matrix script and current PyTorch release matrix behavior. Testing: - D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/configure_jax_release_matrix_test.py github_actions/tests/configure_pytorch_release_matrix_test.py - pre-commit run --files .github/workflows/multi_arch_release.yml .github/workflows/multi_arch_release_linux.yml .github/workflows/multi_arch_release_linux_jax_wheels.yml build_tools/github_actions/configure_jax_release_matrix.py build_tools/github_actions/configure_pytorch_release_matrix.py build_tools/github_actions/tests/configure_jax_release_matrix_test.py build_tools/github_actions/tests/configure_pytorch_release_matrix_test.py Assisted-by: Codex

therock-pr-bot · 2026-06-24T21:49:11Z

❌ PR Check — Action Required

Check	Status	Details
🌿 Branch Name	✅ Pass	—
📝 PR Title/Description	❌ Fail	Error: Title does not follow Conventional Commits style. Expected: start with a valid type (feat, fix, docs, …). Desired format: `type(optional-scope): short description`
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	✅ Pass	—
🔎 pre-commit	✅ Pass	—
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

⚠️ 1 policy check(s) failed. Please address the issues above before this PR can be Reviewed.

🚫 Please fix the failed policies

❌ PR Title/Description

The Not ready to Review label was added to this PR. Once all policies pass, the label is removed automatically.

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

therock-pr-bot · 2026-06-24T21:49:14Z

🚫 Please fix the failed policies before requesting reviews.

The following policy checks failed:

❌ PR Title/Description

The Not ready to Review label has been added to this PR.
Once all policies pass, the label will be removed automatically.

ScottTodd · 2026-06-25T22:10:28Z

fyi @erman-gurses + @geomin12 , I have this queued up once some already in-review PRs are merged.

…ional-jax

Avoid exhaustively asserting the expanded default matrix. Check that the default matrix spans multiple Python versions and JAX refs, and that explicit Python versions narrow every generated row. Testing: - D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/configure_jax_release_matrix_test.py - D:/projects/TheRock/.venv/Scripts/python.exe -m pytest github_actions/tests/configure_pytorch_release_matrix_test.py - pre-commit run --files build_tools/github_actions/tests/configure_jax_release_matrix_test.py Assisted-by: Codex

erman-gurses · 2026-06-30T21:06:04Z

Asking to make sure - little confused about this test. Is this not triggered because of the false option? or because of the matlib build failed? https://github.com/ROCm/TheRock/actions/runs/28469583076/job/84382542062

Trigger a dev release with multi_arch_release_linux.yml, setting build_jax to false
https://github.com/ROCm/TheRock/actions/runs/28469583076

ScottTodd · 2026-06-30T21:18:06Z

Yeah that looks like an unrelated build flake :/ I just retriggered the jobs that failed on that test run.

erman-gurses

Looks good now!

## Motivation Progress on ROCm/TheRock#5634 and ROCm/TheRock#6218. This will be used by ROCm/TheRock#6117 which applies what ROCm/TheRock#5727 added for pytorch to jax: an optional `build_jax` toggle and the `python_version` input to build for a single version instead of all versions. ## Test Plan Dev release in TheRock: https://github.com/ROCm/TheRock/actions/runs/28469583076 ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

## Motivation Progress on #5634 and #6218. This builds on #6117 to provide the JAX equivalent to #6082. Now the dynamic matrix of JAX builds to run will be generated at the _start_ of CI and release pipelines during "setup" / "configure CI" and a table will be included in the summary: <img width="750" height="683" alt="image" src="https://github.com/user-attachments/assets/37b21bd3-1906-4d2f-9732-0602fab9776a" /> ## Technical Details The new `build_jax` option now has these settings: Workflow | `build_jax` value | notes -- | -- | -- `multi_arch_ci.yml` | `false` | Matching previous behavior<br>can later be enabled (opt-in, automatic based on files edited, etc.) `multi_arch_ci_asan.yml` | `false` | New behavior due to #6218 `multi_arch_release.yml` | `inputs.build_jax`<br>(default `true`) | Matching previous behavior `multi_arch_release_asan.yml` | `false` | New behavior due to #6218 ## Test Plan * Dev release triggered to observe CI configuration: https://github.com/ROCm/TheRock/actions/runs/28470132517 ## Submission Checklist - [x] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

ScottTodd added the ci:skip Skip all CI builds/tests for this PR label Jun 24, 2026

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

This was referenced Jun 26, 2026

Use configure_pytorch_release_matrix.py to drive the CI job matrix #6082

Merged

fix(jax): skip JAX wheel builds/tests on ASan builds #6213

Closed

[ASAN]: Skip building python packages and frameworks in asan pipelines until supported #6218

Open

ScottTodd added 2 commits June 30, 2026 11:23

Merge remote-tracking branch 'upstream/main' into users/scotttodd/opt…

f99c177

…ional-jax

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

Fix bad merge

99c2d48

ScottTodd mentioned this pull request Jun 30, 2026

Add python_version input to jax release workflow ROCm/rockrel#70

Merged

1 task

ScottTodd marked this pull request as ready for review June 30, 2026 20:16

ScottTodd requested review from erman-gurses and geomin12 June 30, 2026 20:16

ScottTodd mentioned this pull request Jun 30, 2026

Thread JAX build matrix through setup #6270

Merged

1 task

ScottTodd mentioned this pull request Jun 30, 2026

Test generated workflow matrix schemas #6277

Open

1 task

erman-gurses approved these changes Jun 30, 2026

View reviewed changes

ScottTodd merged commit 4477997 into main Jun 30, 2026
155 of 170 checks passed

ScottTodd deleted the users/scotttodd/optional-jax branch June 30, 2026 23:44

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add optional JAX release build toggle#6117

Add optional JAX release build toggle#6117
ScottTodd merged 4 commits into
mainfrom
users/scotttodd/optional-jax

ScottTodd commented Jun 24, 2026 •

edited

Loading

Uh oh!

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

edited

Loading

Uh oh!

therock-pr-bot Bot commented Jun 24, 2026

Uh oh!

ScottTodd commented Jun 25, 2026

Uh oh!

erman-gurses commented Jun 30, 2026 •

edited

Loading

Uh oh!

ScottTodd commented Jun 30, 2026

Uh oh!

erman-gurses left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ScottTodd commented Jun 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Submission Checklist

Uh oh!

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

Uh oh!

❌ PR Check — Action Required

Uh oh!

therock-pr-bot Bot commented Jun 24, 2026

Uh oh!

ScottTodd commented Jun 25, 2026

Uh oh!

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

Uh oh!

Uh oh!

ScottTodd commented Jun 30, 2026

Uh oh!

erman-gurses left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ScottTodd commented Jun 24, 2026 •

edited

Loading

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

edited

Loading

erman-gurses commented Jun 30, 2026 •

edited

Loading