Enable neural network tests on ROCm

[phambinhfin]

Enable NN tests to run on ROCm by using XLA implementation instead of cuDNN (which is NVIDIA-only) and fixing compute capability skip checks.

Changes:

• testScaledMatmul: skip compute capability check on ROCm (works on ROCm)
• testScaledDotGeneral: skip compute capability check on ROCm (works on ROCm)
• testDotProductAttention: add ROCm skip for cuDNN impl (XLA impl still runs)
• testDotProductAttentionMask: use XLA instead of cuDNN on ROCm
• testDotProductAttentionBiasGradient: use XLA instead of cuDNN on ROCm

Related commits from ROCm/jax PRs #604, #637, #640

[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 significantly enhances the compatibility of neural network tests with ROCm devices. It achieves this by intelligently adjusting compute capability checks to accommodate ROCm environments and by strategically replacing NVIDIA-specific cuDNN implementations with the more general XLA implementation for several key attention-related tests. The overall impact is a broader test coverage and improved functionality of JAX on AMD GPUs.

Highlights

• ROCm Compatibility for Scaled Operations: The testScaledMatmul and testScaledDotGeneral tests now bypass CUDA compute capability checks when running on ROCm devices, allowing these tests to execute successfully on AMD GPUs.
• Conditional cuDNN Usage for Attention Tests: The testDotProductAttention test now explicitly skips the cuDNN implementation on ROCm, ensuring that the XLA implementation is used instead, as cuDNN is NVIDIA-specific.
• XLA Implementation for Dot Product Attention Mask and Bias Gradient: The testDotProductAttentionMask and testDotProductAttentionBiasGradient tests have been modified to default to the XLA implementation on ROCm devices, moving away from the NVIDIA-specific cuDNN implementation.
• Refined Backend Checks: Assertions related to checking the cuDNN backend are now conditional, only executing when cuDNN is actually in use (i.e., not on ROCm devices).

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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

The pull request successfully enables neural network tests on ROCm by conditionally using the XLA implementation instead of cuDNN and adjusting compute capability skip checks. The changes are well-aligned with the pull request description, ensuring that tests run correctly on ROCm devices where cuDNN is not supported or desired. The modifications are consistent across the affected test functions, correctly isolating CUDA-specific logic and enabling XLA as a fallback. No issues were found in the reviewed changes.