[ROCm] Add ROCm backward compatibility test for lu_pivots_to_permutation

[AratiGanesh]

Add ROCm backward compatibility test for lu_pivots_to_permutation

This PR adds a backward compatibility test for the LU pivots to permutation operation on ROCm platform, ensuring that serialized MLIR modules from older JAX versions can still be correctly loaded and executed.

Changes:

• Add test_rocm_lu_pivots_to_permutation test method
• Add rocm_lu_pivots_to_permutation test data (serialized Feb 4, 2026)
• Test covers hip_lu_pivots_to_permutation custom call

[gemini-code-assist]

Summary of Changes

Hello @AratiGanesh, 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 enhances the JAX testing infrastructure by introducing a backward compatibility test for the lu_pivots_to_permutation operation on ROCm. The primary goal is to guarantee that MLIR modules, which were serialized from previous JAX versions and utilize this specific operation, can continue to be loaded and executed correctly on the ROCm platform, thereby preventing regressions and ensuring long-term stability.

Highlights

• ROCm Backward Compatibility Test: Added a new backward compatibility test for the lu_pivots_to_permutation operation specifically for the ROCm platform.
• New Test Data: Introduced new test data for rocm_lu_pivots_to_permutation, serialized on February 4, 2026, to ensure future compatibility.
• Custom Call Coverage: The new test covers the hip_lu_pivots_to_permutation custom call, verifying its behavior with older serialized MLIR modules.

🧠 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.

Changelog

• _src/internal_test_util/export_back_compat_test_data/rocm_lu_pivots_to_permutation.py
  • Added a new Python file containing serialized MLIR module text, expected outputs, and metadata for the ROCm lu_pivots_to_permutation backward compatibility test.
• tests/export_back_compat_test.py
  • Imported the new rocm_lu_pivots_to_permutation test data module.
  • Integrated the rocm_lu_pivots_to_permutation test data into the test_custom_call_coverage method.
  • Added a new test method, test_rocm_lu_pivots_to_permutation, to execute the backward compatibility test for the ROCm platform.

Activity

• The pull request was created by AratiGanesh to add a ROCm backward compatibility test for lu_pivots_to_permutation.
• A new test data file was added, and the main test suite was updated to include and run this new test.

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 adds a backward compatibility test for lu_pivots_to_permutation on the ROCm platform. The changes are generally good, adding the necessary test data and test case. I have two suggestions to improve the code. First, to make the custom call coverage test more robust, the now-covered hip_lu_pivots_to_permutation should be removed from the list of manually ignored targets. Second, the new test for ROCm is very similar to the existing CUDA test; they could be combined into a single parameterized test to reduce duplication and improve maintainability.

[gnecula]

Thanks!

[gnecula]

There seem to be some related errors

[gnecula]

Please rebase. In general, because you have a series of PRs that modify the same file, it is easier if we try to submit them one by one, with rebasing in the middle. I sometimes can do that myself internally in Google, but it involves more steps.