Skip to content

feat: Add JAX release support v0.10.2#6195

Merged
magaonka-amd merged 1 commit into
mainfrom
users/magaonka-amd/add-jax-v0.10.2
Jun 29, 2026
Merged

feat: Add JAX release support v0.10.2#6195
magaonka-amd merged 1 commit into
mainfrom
users/magaonka-amd/add-jax-v0.10.2

Conversation

@magaonka-amd

@magaonka-amd magaonka-amd commented Jun 28, 2026

Copy link
Copy Markdown
Contributor

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

- Add rocm-jaxlib-v0.10.2 to the multi-arch release matrix, wired
  identically to v0.10.0 (ROCm/jax, manylinux build_mode, device-all).
- Use the default XLA allocator (XLA_PYTHON_CLIENT_ALLOCATOR=default)
  only for the v0.10.2 test cell; all other releases keep platform.
- Document 0.10.2 in the JAX support table.
@therock-pr-bot

therock-pr-bot Bot commented Jun 28, 2026

Copy link
Copy Markdown

✅ All Checks Passed — Ready for Review

Check Status Details
🌿 Branch Name ✅ Pass
📝 PR Title/Description ✅ Pass
Forbidden Files ✅ Pass
🧪 Unit Test ✅ Pass PR does not contain code files — Unit Test auto-passed
🔎 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

Copy link
Copy Markdown

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

@erman-gurses erman-gurses left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

LGTM if rocm-jaxlib-v0.10.2 is fully ready - @tsrw2048?

@tsrw2048

Copy link
Copy Markdown

LGTM if rocm-jaxlib-v0.10.2 is fully ready - @tsrw2048?

@erman-gurses - Yes, @magaonka-amd has been working on 0.10.2 in parallel.

@magaonka-amd magaonka-amd merged commit 939b170 into main Jun 29, 2026
79 of 91 checks passed
@magaonka-amd magaonka-amd deleted the users/magaonka-amd/add-jax-v0.10.2 branch June 29, 2026 14:16
@github-project-automation github-project-automation Bot moved this from TODO to Done in TheRock Triage Jun 29, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

3 participants