Skip to content

Add support for Jax v.0.9.0 #3717

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JeniferC99 merged 8 commits into from
Mar 9, 2026
Merged

Add support for Jax v.0.9.0 #3717

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
b08b587
Change JAX version to 0.9.0 (rocm-jaxlib-v0.9.0)
Feb 25, 2026
2109196
Conditionally add --jax-source-dir only for jax_ref 0.8.2
Feb 27, 2026
4d082b3
add jax path for 0.9.0 too.
Feb 27, 2026
6ee5c16
add other jax package version
Mar 3, 2026
83bfdc9
Merge branch 'main' into users/jechrist/jax0.9.0
JeniferC99 Mar 3, 2026
1dd3a69
Merge branch 'main' into users/jechrist/jax0.9.0
JeniferC99 Mar 6, 2026
9adefde
Merge branch 'main' into users/jechrist/jax0.9.0
JeniferC99 Mar 9, 2026
691f47f
Merge branch 'main' into users/jechrist/jax0.9.0
JeniferC99 Mar 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion .github/workflows/build_linux_jax_wheels.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ jobs:
run: |
ls -lah
pushd rocm-jax
[[ "${{ inputs.jax_ref }}" == *"0.8.2"* ]] && SOURCE_ARG="--jax-source-dir=${{ github.workspace }}/jax" || SOURCE_ARG=""
[[ "${{ inputs.jax_ref }}" != *"0.8.0"* ]] && SOURCE_ARG="--jax-source-dir=${{ github.workspace }}/jax" || SOURCE_ARG=""
Copy link
Copy Markdown
Contributor

@charleshofer charleshofer Mar 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: Might want to make this a less than comparison so that we exclude 0.7.x as well. But also probably not a huge deal since we're not building with 0.7.x.

python3 build/ci_build \
--compiler=clang \
--python-versions="${{ inputs.python_version }}" \
Expand Down Expand Up @@ -256,6 +256,8 @@ jobs:
jax_ref: ${{ needs.build_jax_wheels.outputs.jax_ref }}
jax_version: ${{ needs.build_jax_wheels.outputs.jax_version }}
jaxlib_version: ${{ needs.build_jax_wheels.outputs.jaxlib_version }}
jax_plugin_version: ${{ needs.build_jax_wheels.outputs.jax_plugin_version }}
jax_pjrt_version: ${{ needs.build_jax_wheels.outputs.jax_pjrt_version }}
test_runs_on: ${{ needs.generate_target_to_run.outputs.test_runs_on }}

upload_jax_wheels:
Expand Down
2 changes: 1 addition & 1 deletion .github/workflows/release_portable_linux_jax_wheels.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ jobs:
fail-fast: false
matrix:
python_version: ["3.11", "3.12", "3.13", "3.14"]
jax_ref: ["rocm-jaxlib-v0.8.0", "rocm-jaxlib-v0.8.2"]
jax_ref: ["rocm-jaxlib-v0.8.0", "rocm-jaxlib-v0.8.2", "rocm-jaxlib-v0.9.0"]

uses: ./.github/workflows/build_linux_jax_wheels.yml
with:
Expand Down
28 changes: 24 additions & 4 deletions .github/workflows/test_linux_jax_wheels.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,15 @@ on:
required: false
type: string
jaxlib_version:
description: "Full jaxlib version including rocm suffix (e.g. 0.8.0+rocm7.12.0). Extracted from built wheels by write_jax_versions.py."
description: "jaxlib wheel version (e.g. 0.9.0+rocm7 or 0.8.0+rocm7.12.0). Extracted from built wheels by write_jax_versions.py."
required: false
type: string
jax_plugin_version:
description: "jax_rocm7_plugin wheel version. Extracted from built wheels by write_jax_versions.py."
required: false
type: string
jax_pjrt_version:
description: "jax_rocm7_pjrt wheel version. Extracted from built wheels by write_jax_versions.py."
required: false
type: string
ref:
Expand Down Expand Up @@ -103,7 +111,15 @@ on:
required: false
type: string
jaxlib_version:
description: "Full jaxlib version with rocm suffix (e.g. 0.8.0+rocm7.12.0). Leave empty to auto-compute from rocm_version."
description: "jaxlib wheel version (e.g. 0.9.0+rocm7 or 0.8.0+rocm7.12.0). Leave empty to auto-compute from rocm_version."
required: false
type: string
jax_plugin_version:
description: "jax_rocm7_plugin wheel version. Leave empty to use jaxlib_version."
required: false
type: string
jax_pjrt_version:
description: "jax_rocm7_pjrt wheel version. Leave empty to use jaxlib_version."
required: false
type: string
test_runs_on:
Expand Down Expand Up @@ -206,6 +222,8 @@ jobs:
echo "Using versions passed from build workflow"
echo "JAXLIB_VERSION=${{ inputs.jaxlib_version }}" >> "$GITHUB_ENV"
echo "JAX_VERSION=${{ inputs.jax_version }}" >> "$GITHUB_ENV"
echo "JAX_PLUGIN_VERSION=${{ inputs.jax_plugin_version }}" >> "$GITHUB_ENV"
echo "JAX_PJRT_VERSION=${{ inputs.jax_pjrt_version }}" >> "$GITHUB_ENV"
else
echo "Computing versions from rocm_version and rocm-jax requirements"
pip install packaging
Expand All @@ -218,12 +236,14 @@ jobs:
run: |
echo "Installing JAX packages:"
echo " JAXLIB_VERSION=${JAXLIB_VERSION}"
echo " JAX_PLUGIN_VERSION=${JAX_PLUGIN_VERSION}"
echo " JAX_PJRT_VERSION=${JAX_PJRT_VERSION}"
echo " JAX_VERSION=${JAX_VERSION}"
# Install jaxlib/plugin/pjrt from the GPU-family index; install jax from PyPI to match the version
pip install --index-url "${{ env.WHEEL_INDEX_URL }}" \
"jaxlib==${JAXLIB_VERSION}" \
"jax_rocm7_plugin==${JAXLIB_VERSION}" \
"jax_rocm7_pjrt==${JAXLIB_VERSION}"
"jax_rocm7_plugin==${JAX_PLUGIN_VERSION}" \
"jax_rocm7_pjrt==${JAX_PJRT_VERSION}"
pip install jax==${JAX_VERSION}

- name: Run JAX tests
Expand Down
Loading