jax wheel tests by kiran-thumma · Pull Request #2247 · ROCm/TheRock

kiran-thumma · 2025-11-21T06:58:36Z

Motivation

Add an automated GPU validation test that verifies built JAX wheel (.whl) files are functional.
We currently lack automated wheel-level GPU validation. This test will catch regressions in packaging, compatibility issues, missing binaries, or other wheel-level failures before wheels are promoted.
Faster feedback, reduces risk of shipping broken GPU wheels.

Technical Details

New workflow: test_linux_jax_wheels.yml
Test type: Core , utils, scipy stats and multi device tests
Purpose: Install built JAX wheels and run tests.
Diagram:

Test Plan

Local Verification
CI Verification

Test Result

Workflow run: Workflow run: https://github.com/ROCm/TheRock/actions/runs/19603110752/job/56139806408

Submission Checklist

All tests pass on at least one GPU configuration in CI
Workflow has a manual dispatch input for testing arbitrary wheel URLs
Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

kiran-thumma · 2025-11-25T17:21:02Z

Latest: https://github.com/ROCm/TheRock/actions/runs/19676348633/job/56364211047

geomin12 · 2025-11-25T17:57:50Z

+          # Extract JAX version from requirements.txt (e.g., "jax==0.8.0")
+          JAX_VERSION=$(tr -d ' ' < jax/build/requirements.txt \
+          | grep -E '^(jax|jaxlib)==[^#]+' | head -n1 | cut -d'=' -f3)
+          echo "JAX_VERSION=$JAX_VERSION" >> "$GITHUB_ENV"


could we instead pull it from inputs.jax_ref ?

No, The exact version of each can only be retrieved from the requirements.txt. As the plan to support both master and a release branch. We will need the jax version from requirements.txt

could you please add a bit more context to this?

you build jax with version x.y. but inputs.jax_ref is only used to get the jax checkout to fetch the correct tests?
how would the version look differently between master and test?
how do you then test a new jax version?

it would be also good not to use bash but some python code to exract the version. does something like

import jax print(jax.__version__)

work?

agree with Laura, please add docs for this or a comment

as i was reading this, i am confused on why (also, there is complex bash too)

Sure, I will add comments. And wont work as we need the value prior to installing jax.
import jax
print(jax.__version__)

Yes let's rework this step (can be in a follow-up). There is too much bash magic here independent of the workflow inputs. We probably want to support building multiple jax versions (at least "stable" and "latest"), so we can't just use the version in a pinned requirements file.

kiran-thumma · 2025-11-26T08:57:47Z

Latest Run : https://github.com/ROCm/TheRock/actions/runs/19695651948

HereThereBeDragons

thanks so far for removing the docker image!

see below some questions and comments i have

HereThereBeDragons · 2025-11-26T09:55:25Z

+          # Extract JAX version from requirements.txt (e.g., "jax==0.8.0")
+          JAX_VERSION=$(tr -d ' ' < jax/build/requirements.txt \
+          | grep -E '^(jax|jaxlib)==[^#]+' | head -n1 | cut -d'=' -f3)
+          echo "JAX_VERSION=$JAX_VERSION" >> "$GITHUB_ENV"


could you please add a bit more context to this?

you build jax with version x.y. but inputs.jax_ref is only used to get the jax checkout to fetch the correct tests?
how would the version look differently between master and test?
how do you then test a new jax version?

it would be also good not to use bash but some python code to exract the version. does something like

import jax print(jax.__version__)

work?

geomin12 · 2025-11-26T17:26:14Z

+          # Extract JAX version from requirements.txt (e.g., "jax==0.8.0")
+          JAX_VERSION=$(tr -d ' ' < jax/build/requirements.txt \
+          | grep -E '^(jax|jaxlib)==[^#]+' | head -n1 | cut -d'=' -f3)
+          echo "JAX_VERSION=$JAX_VERSION" >> "$GITHUB_ENV"


agree with Laura, please add docs for this or a comment

as i was reading this, i am confused on why (also, there is complex bash too)

ScottTodd · 2025-11-26T19:30:16Z

Please add a title for this PR, not "Users/kithumma/jax test". Also tag relevant issues tracking the work, like #1985.

geomin12

looking good! a few clarifications are needed for my understanding and some housekeeping items, but aside from that, looks good

geomin12

looking good! just some concerns about the inputs.ref as well as missing inputs in the test file

kiran-thumma · 2025-12-02T16:49:37Z

@geomin12 Created an issue 2383 which is to refactor all the workflows with inputs.repository, inputs.ref and github.repository. Please add your suggestions in that issue what do you need instead of this logic.

geomin12 · 2025-12-02T16:57:21Z

@geomin12 Created an issue 2383 which is to refactor all the workflows with inputs.repository, inputs.ref and github.repository. Please add your suggestions in that issue what do you need instead of this logic.

that's fine to do cleanup in another PR, but my concern is this code won't work based on what I am seeing

My first concern is the workflow_call repository and ref inputs. in the "test" workflow (that depends on build), ref is passed but repository is not passed. this will cause an error if a ref is passed but doesn't known which repository it belongs too
{{ inputs.ref || '' }} will also cause an error. what if inputs ref is empty? an empty string ref is invalid and will also error
repository and ref will default to ROCm/TheRock and whichever branch it is on during a workflow_call, which makes those two inputs unnecessary (i guess why add unnecessary inputs)

points 1 and 2 also contradict with point 3, which is why I'm curious on why we add those two inputs to workflow_call when they will always default (no point in extra unnecessary logic)

geomin12

looks good, i just see one bug then should be good to go

ScottTodd · 2025-12-02T20:52:54Z

+      - name: Run JAX tests
+        run: |
+            pytest jax/jax_tests/tests/multi_device_test.py -q --log-cli-level=INFO
+            pytest jax/jax_tests/tests/core_test.py -q --log-cli-level=INFO
+            pytest jax/jax_tests/tests/util_test.py -q --log-cli-level=INFO
+            pytest jax/jax_tests/tests/scipy_stats_test.py -q --log-cli-level=INFO


This will exit on the first test failure, when we probably want to keep going and report all failures together. Fine to improve in a follow-up.

ScottTodd · 2025-12-02T20:54:41Z

+          # Extract JAX version from requirements.txt (e.g., "jax==0.8.0")
+          JAX_VERSION=$(tr -d ' ' < jax/build/requirements.txt \
+          | grep -E '^(jax|jaxlib)==[^#]+' | head -n1 | cut -d'=' -f3)
+          echo "JAX_VERSION=$JAX_VERSION" >> "$GITHUB_ENV"


Yes let's rework this step (can be in a follow-up). There is too much bash magic here independent of the workflow inputs. We probably want to support building multiple jax versions (at least "stable" and "latest"), so we can't just use the version in a pinned requirements file.

geomin12

looks good, just one fix! this will cause errors in test jax file

geomin12

sweet, thanks for the updates! you can also close that other issue as the repository and ref is needed since this will be a part of TheRock releases

ScottTodd

A few small comments but otherwise LGTM

kiran-thumma · 2025-12-04T17:17:55Z

A few small comments but otherwise LGTM

Updated all as suggested.

Got approval from 2/3 reviewers

## Motivation - Add an automated GPU validation test that verifies built JAX wheel (.whl) files are functional. - We currently lack automated wheel-level GPU validation. This test will catch regressions in packaging, compatibility issues, missing binaries, or other wheel-level failures before wheels are promoted. - Faster feedback, reduces risk of shipping broken GPU wheels.  ## Technical Details - New workflow: test_linux_jax_wheels.yml - Test type: Core , utils, scipy stats and multi device tests - Purpose: Install built JAX wheels and run tests. - Diagram: <img width="170" height="359" alt="Screenshot 2025-11-23 140535" src="https://github.com/user-attachments/assets/f7b757f3-7013-4253-ae5b-a3a12934e6b6" />  ## Test Plan - Local Verification - CI Verification  ## Test Result - Workflow run: Workflow run: https://github.com/ROCm/TheRock/actions/runs/19603110752/job/56139806408  ## Submission Checklist - [x] All tests pass on at least one GPU configuration in CI - [x] Workflow has a manual dispatch input for testing arbitrary wheel URLs - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

github-project-automation Bot added this to TheRock Triage Nov 21, 2025

github-project-automation Bot moved this to TODO in TheRock Triage Nov 21, 2025

kiran-thumma requested review from HereThereBeDragons and geomin12 November 25, 2025 17:27

kiran-thumma marked this pull request as ready for review November 25, 2025 17:28

geomin12 reviewed Nov 25, 2025

View reviewed changes

This was referenced Nov 26, 2025

add jax core test #2035

Closed

READme update for jax #2135

Merged

kiran-thumma requested a review from geomin12 November 26, 2025 09:12

HereThereBeDragons previously requested changes Nov 26, 2025

View reviewed changes

kiran-thumma requested a review from HereThereBeDragons November 26, 2025 10:49

geomin12 reviewed Nov 26, 2025

View reviewed changes

ScottTodd mentioned this pull request Nov 27, 2025

Add actionlint as a pre-commit hook #2317

Merged

kiran-thumma changed the title ~~Users/kithumma/jax test~~ jax wheel tests Nov 28, 2025

kiran-thumma requested a review from geomin12 November 28, 2025 22:28

geomin12 reviewed Dec 1, 2025

View reviewed changes

kiran-thumma requested a review from geomin12 December 1, 2025 22:23

geomin12 reviewed Dec 2, 2025

View reviewed changes

kiran-thumma requested a review from geomin12 December 2, 2025 16:50

geomin12 reviewed Dec 2, 2025

View reviewed changes

Comment thread .github/workflows/build_linux_jax_wheels.yml

Comment thread .github/workflows/test_linux_jax_wheels.yml

ScottTodd reviewed Dec 2, 2025

View reviewed changes

kiran-thumma added 4 commits December 3, 2025 00:40

jax test

797fb61

jax test

e958aeb

file name update

ca43e07

update workflow dispatch

8699039

kiran-thumma added 9 commits December 3, 2025 00:47

update

e1d7d04

fixing merge conflicts

3a4f914

update

744f1e3

remove pip upgrade and add in python setup

50cf5b9

remove symlinks

5e820c9

udpates

9987b72

updated version to exat version

9c15b8c

updated numpy version

41c4d9b

fixing merge conflicts

537526c

kiran-thumma force-pushed the users/kithumma/jax-test branch from 8e96f55 to 537526c Compare December 3, 2025 00:50

kiran-thumma added 4 commits December 3, 2025 01:24

update python_version parameter

0dc4566

update repository workflow_call

5d5f405

udpate version scipy, ml_dtypes, numpy

ce3e0e9

update amdgpu_family

8fa37f8

kiran-thumma requested review from ScottTodd and geomin12 December 3, 2025 08:37

geomin12 requested changes Dec 3, 2025

View reviewed changes

Comment thread .github/workflows/test_jax_dockerfile.yml

update merge conflicts

df44c8c

kiran-thumma requested a review from geomin12 December 3, 2025 20:22

geomin12 approved these changes Dec 3, 2025

View reviewed changes

ScottTodd reviewed Dec 4, 2025

View reviewed changes

Comment thread .github/workflows/build_linux_jax_wheels.yml Outdated

Comment thread .github/workflows/test_jax_dockerfile.yml Outdated

Comment thread .github/workflows/test_linux_jax_wheels.yml

updates as per comments

814582b

kiran-thumma requested a review from ScottTodd December 4, 2025 17:18

ScottTodd approved these changes Dec 4, 2025

View reviewed changes

kiran-thumma merged commit d53002e into main Dec 4, 2025
9 checks passed

github-project-automation Bot moved this from TODO to Done in TheRock Triage Dec 4, 2025

kiran-thumma deleted the users/kithumma/jax-test branch December 4, 2025 17:33

Conversation

kiran-thumma commented Nov 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

kiran-thumma commented Nov 25, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

kiran-thumma commented Nov 26, 2025

Uh oh!

HereThereBeDragons left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ScottTodd commented Nov 26, 2025

Uh oh!

geomin12 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

geomin12 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kiran-thumma commented Dec 2, 2025

kiran-thumma commented Nov 21, 2025 •

edited

Loading

geomin12 commented Dec 2, 2025 •

edited

Loading

geomin12 left a comment •

edited

Loading