Fix ROCm lib names copy pattern to use wildcard version

[phambinhfin]

The build script failed to copy ROCm 7 wheels due to hardcoded version
pattern. Changed to use wildcard (rocm*) to match any ROCm version.
Changes:
Use rocm* wildcard instead of rocm{version} in wheel pattern
Add jax-rocm-pjrt to py-tagged wheels list
Supports jax_rocm60_, jax_rocm7_, and future versions

[gemini-code-assist]

Summary of Changes

Hello @phambinhfin, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical issue in the build script where ROCm 7 wheels were not being copied due to a rigid version matching pattern. By introducing a wildcard ("rocm*") for ROCm versions, the script can now correctly identify and copy wheels for current and future ROCm releases. Additionally, it ensures that "jax-rocm-pjrt" wheels are properly categorized with a pure Python tag.

Highlights

• ROCm Wheel Copy Pattern: Updated the build script to use a "rocm*" wildcard pattern instead of hardcoded "rocm{version}" for copying ROCm library names, resolving issues with ROCm 7 wheels and supporting future versions.
• jax-rocm-pjrt Python Tagging: Added "jax-rocm-pjrt" to the list of wheels that are tagged as pure Python ("py"), ensuring correct packaging and distribution.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request fixes an issue with building ROCm wheels by using a wildcard to match different ROCm versions, which makes the build script more flexible. The change to add jax-rocm-pjrt to the list of py-tagged wheels is also correct.

My review includes one point of feedback on the build/build.py file. I've identified a potential issue where using a wildcard in the wheel_dir variable could cause problems for future editable builds for ROCm. While it doesn't break current functionality, it introduces fragility into the script. I've suggested making the code more robust to handle this case.

[phambinhfin]

/gemini review

[gemini-code-assist]

Code Review

This pull request addresses an issue with building ROCm wheels for different versions by using a wildcard pattern. The changes look good and correctly implement the intended fix. Specifically, the build script now uses rocm* to match any ROCm version for non-editable builds, which makes the process more flexible and future-proof. Additionally, jax-rocm-pjrt is now correctly tagged as a Python-version-independent wheel. The changes are clear and well-contained.