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:

From Multi-Arch Release Linux JAX Wheels workflow: https://github.com/ROCm/rockrel/actions/runs/28078871973 (Completed)
From Multi-Arch Release workflow https://github.com/ROCm/rockrel/actions/runs/28078897916 (Completed)
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 Checks Passed — Ready for Review

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

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

📖 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-24T16:34:33Z

@ScottTodd , Is there any blocker for now from your side? If not. good to merge?

ScottTodd · 2026-06-24T16:35:30Z

I can do another review today. 15 other PRs in the queue.

ScottTodd · 2026-06-24T17:52:54Z

+      build_mode:
+        description: "Build mode: native or manylinux."
+        type: string
+        required: true


nit: If there are only two expected string values, for workflow_dispatch you could use a "choice" input type to make this easier to use correctly.

https://docs.github.com/en/actions/reference/workflows-and-actions/workflow-syntax#onworkflow_dispatchinputs

ScottTodd · 2026-06-24T17:53:54Z

+      jax_repository:
+        description: "Repository containing the JAX release branch."
+        type: string
+        required: true


Can you add a default here, or at least mention which repositories are expected in the description? What syntax is this expecting, ROCm/jax or https://github.com/ROCm/jax?

ScottTodd · 2026-06-24T17:54:42Z

+      gfx_arch:
+        description: "GFX architecture used for the manylinux image build."
+        type: string
+        required: false
+        default: ""


Add an example to the description? We have a few syntax styles (gfx94X, gfx942, gfx942-dcgpu, etc.) and that helps a lot when working across multiple workflows.

ScottTodd · 2026-06-24T17:58:00Z

      repository:
        description: "Repository to checkout. Defaults to github.repository."
        type: string
        default: ""
      ref:
        description: "Branch, tag, or SHA to checkout. Defaults to the triggering ref."
        type: string
        default: ""


These repository and ref inputs aren't commonly provided when triggered by developers, so I'd prefer to keep them as the last inputs in the list. Looks like a few other workflows have similarly started adding inputs beyond them though.

I would move jax_repository, build_mode, and gfx_arch near the top of the inputs list since they will be set frequently.

ScottTodd · 2026-06-24T18:01:01Z

+      - name: Build JAX Wheels in manylinux container
+        if: ${{ inputs.build_mode == 'manylinux' }}
+        working-directory: jax-source
+        env:
+          ROCM_VERSION: ${{ inputs.rocm_version }}
+          PYTHON_VERSION: ${{ inputs.python_version }}
+        run: |
+          docker run --rm \


Can you also update the docs at https://github.com/ROCm/TheRock/tree/main/external-builds/jax#build-instructions ? (The general progression for development work and feature enablement should be local builds --> dev builds --> release builds, this jumps straight to release builds in github actions workflows which is the hardest to debug)

ScottTodd · 2026-06-24T18:03:46Z

+      rocm_package_find_links_url:
+        description: "ROCm package index / find-links URL for the JAX manylinux build."
+        type: string
+        default: "https://rocm.devreleases.amd.com/whl-multi-arch/"


Remove default that won't work here too

ScottTodd · 2026-06-24T18:05:02Z

+      - name: Install ROCm packages
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install --pre \
+            --index-url "${ROCM_PACKAGE_FIND_LINKS_URL}" \


Wait. Which URL format do you want? They are different. The variable name here says "find links url" but it's installing with --index-url?

ScottTodd · 2026-06-24T18:06:48Z

+            --index-url "${ROCM_PACKAGE_FIND_LINKS_URL}" \
+            rocm \
+            rocm-sdk-devel \
+            rocm-sdk-device-gfx942


test_amdgpu_family is an input but this assumes that gfx942 is all that needs to be installed?

Use this code instead:

ROCm:

TheRock/.github/workflows/test_rocm_wheels.yml

Lines 147 to 176 in 1371136

- name: Compute ROCm device extras

if: ${{ inputs.kpack_split == 'true' }}

id: device_extras

run: |

python build_tools/github_actions/expand_amdgpu_families.py \

--amdgpu-families "${{ inputs.amdgpu_family }}" \

--output-mode=device-extras

# Key off install_libraries success (not test_libraries) so devel install/test still run

- name: Install rocm[libraries]

id: install_libraries

run: |

pkg="rocm[libraries]"

if [ -n "${{ steps.device_extras.outputs.device_extras }}" ]; then

pkg="${pkg%]},${{ steps.device_extras.outputs.device_extras }}]"

fi

if [ "${{ inputs.kpack_split }}" = "true" ]; then

python build_tools/setup_venv.py ${VENV_DIR} \

--packages "${pkg}==${{ inputs.rocm_version }}" \

--index-url=${{ inputs.package_index_url_base }} \

--find-links=${{ inputs.package_find_links_url }} \

--activate-in-future-github-actions-steps

else

python build_tools/setup_venv.py ${VENV_DIR} \

--packages "${pkg}==${{ inputs.rocm_version }}" \

--index-url=${{ inputs.package_index_url_base }} \

--index-subdir=${{ inputs.amdgpu_family }} \

--find-links=${{ inputs.package_find_links_url }} \

--activate-in-future-github-actions-steps

fi

torch:

TheRock/.github/workflows/test_pytorch_wheels.yml

Lines 148 to 157 in 1371136

# Multi-arch packages use pip extras (e.g. torch[device-gfx942]) to

# select GPU-specific device packages. Expand the family to get the

# comma-separated device extras string for pip install.

- name: Expand GPU family to device extras

if: ${{ inputs.multi_arch }}

id: expand

run: |

python build_tools/github_actions/expand_amdgpu_families.py \

--amdgpu-families "${{ inputs.amdgpu_family }}" \

--output-mode=device-extras

TheRock/.github/workflows/test_pytorch_wheels.yml

Lines 198 to 212 in 1371136

- name: Set up virtual environment

timeout-minutes: 15

run: |

if [ "${{ inputs.multi_arch }}" = "true" ]; then

python build_tools/setup_venv.py ${VENV_DIR} \

--packages "torch[${{ steps.expand.outputs.device_extras }}]==${{ inputs.torch_version }}" \

--index-url=${{ inputs.package_index_url }} \

--activate-in-future-github-actions-steps

else

python build_tools/setup_venv.py ${VENV_DIR} \

--packages torch==${{ inputs.torch_version }} \

--index-url=${{ inputs.package_index_url }} \

--index-subdir=${{ inputs.amdgpu_family }} \

--activate-in-future-github-actions-steps

fi

ScottTodd · 2026-06-24T18:08:22Z

+          python -m pip install --pre \
+            --index-url "${ROCM_PACKAGE_FIND_LINKS_URL}" \
+            rocm \
+            rocm-sdk-devel \


Do JAX tests need rocm-sdk-devel? They shouldn't. Devel packages are for building, and most users of frameworks are not expected to install them.

See my other comment, this should just install rocm[libraries,device-gfxXXXX], or maybe something like jax[device-gfxXXXX] in the future if we want to go the same route as we did for torch extras.

ScottTodd · 2026-06-24T18:11:05Z

+            rocm-sdk-device-gfx942
+          rocm-sdk init
+          echo "ROCM_ROOT=$(rocm-sdk path --root)" >> "$GITHUB_ENV"
+          ln -sfn "$(rocm-sdk path --root)" /opt/rocm


Please don't symlink rocm python packages to /opt/rocm/.

Framework should not be dependent on the devel package at runtime (e.g. in tests)

Python packages exist in python environments, they should not leak outside to the full system like this

therock-pr-bot · 2026-06-25T04:58:20Z

Pre-commit check failed

⛔ pre-commit failed

Please run locally:

python -m pip install pre-commit
pre-commit install
pre-commit run --all-files --show-diff-on-failure

This repo uses .pre-commit-config.yaml.

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

Open

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

Use requirements-jax.txt

1f13220

gulsumgudukbay approved these changes Jun 24, 2026

View reviewed changes

Test

41c015b

erman-gurses requested a review from ScottTodd June 24, 2026 05:51

Remove condition for determine_version.py

099878f

erman-gurses mentioned this pull request Jun 24, 2026

Pass ROCm package index URL through JAX release workflow ROCm/rockrel#65

Open

ScottTodd reviewed Jun 24, 2026

View reviewed changes

erman-gurses added 6 commits June 24, 2026 18:23

Add options for build_mode

88fc90d

Add default for jax_repository

22c916c

Update gfx_arch description

8b6f2f5

Move repository and ref inputs to the bottom

9d58a73

Move jax_repositor,y build_mode, and gfx_arch inputs up

084898b

Solve --find_links inconsistency use index_url

382f39e

ScottTodd mentioned this pull request Jun 24, 2026

Add optional JAX release build toggle #6117

Draft

1 task

Add package_index_url output for better ROCm package indexing

dcf62ab

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

Add unit test

50c652d

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

erman-gurses added 3 commits June 25, 2026 02:41

Increase shared memory size

9b1e7ed

Reduce memory pressure

ff1f03f

Test memory

f7448a4

erman-gurses and others added 6 commits June 25, 2026 05:06

Test

363b89c

Test

0bdd100

Get original testing script

3d46f1b

Update README.md

c15d36d

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

35a873a

Test GPU memory error

c6bacdd

	- name: Compute ROCm device extras
	if: ${{ inputs.kpack_split == 'true' }}
	id: device_extras
	run: \|
	python build_tools/github_actions/expand_amdgpu_families.py \
	--amdgpu-families "${{ inputs.amdgpu_family }}" \
	--output-mode=device-extras

	# Key off install_libraries success (not test_libraries) so devel install/test still run
	- name: Install rocm[libraries]
	id: install_libraries
	run: \|
	pkg="rocm[libraries]"
	if [ -n "${{ steps.device_extras.outputs.device_extras }}" ]; then
	pkg="${pkg%]},${{ steps.device_extras.outputs.device_extras }}]"
	fi
	if [ "${{ inputs.kpack_split }}" = "true" ]; then
	python build_tools/setup_venv.py ${VENV_DIR} \
	--packages "${pkg}==${{ inputs.rocm_version }}" \
	--index-url=${{ inputs.package_index_url_base }} \
	--find-links=${{ inputs.package_find_links_url }} \
	--activate-in-future-github-actions-steps
	else
	python build_tools/setup_venv.py ${VENV_DIR} \
	--packages "${pkg}==${{ inputs.rocm_version }}" \
	--index-url=${{ inputs.package_index_url_base }} \
	--index-subdir=${{ inputs.amdgpu_family }} \
	--find-links=${{ inputs.package_find_links_url }} \
	--activate-in-future-github-actions-steps
	fi

	# Multi-arch packages use pip extras (e.g. torch[device-gfx942]) to
	# select GPU-specific device packages. Expand the family to get the
	# comma-separated device extras string for pip install.
	- name: Expand GPU family to device extras
	if: ${{ inputs.multi_arch }}
	id: expand
	run: \|
	python build_tools/github_actions/expand_amdgpu_families.py \
	--amdgpu-families "${{ inputs.amdgpu_family }}" \
	--output-mode=device-extras

	- name: Set up virtual environment
	timeout-minutes: 15
	run: \|
	if [ "${{ inputs.multi_arch }}" = "true" ]; then
	python build_tools/setup_venv.py ${VENV_DIR} \
	--packages "torch[${{ steps.expand.outputs.device_extras }}]==${{ inputs.torch_version }}" \
	--index-url=${{ inputs.package_index_url }} \
	--activate-in-future-github-actions-steps
	else
	python build_tools/setup_venv.py ${VENV_DIR} \
	--packages torch==${{ inputs.torch_version }} \
	--index-url=${{ inputs.package_index_url }} \
	--index-subdir=${{ inputs.amdgpu_family }} \
	--activate-in-future-github-actions-steps
	fi

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 Checks Passed — Ready for Review

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 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ScottTodd commented Jun 24, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

therock-pr-bot Bot commented Jun 25, 2026

Pre-commit check failed

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 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 24, 2026 •

edited

Loading