jax
72c052e8 - PR #34862: [ROCm] Added back compatibility test for hipsolver_geqrf

Commit
12 days ago
PR #34862: [ROCm] Added back compatibility test for hipsolver_geqrf Imported from GitHub PR https://github.com/jax-ml/jax/pull/34862 ## Motivation The back compatibility test for GEQRF was skipped on ROCm devices because there was no back compatibility test data file. The data file has been added and is now referenced in the back compatibility unit test: `test_gpu_qr_solver_geqrf` allowing it to pass. Consequently, the FFI call for ORGQR is also tested so both GEQRF and ORGQR no longer need to be manually added to the covered targets list and have been removed. ## Technical Details - Added test data file `rocm_qr_hipsolver_geqrf.py`. - Modified existing unit test `test_gpu_qr_solver_geqrf` to support ROCm devices. - Removed `hipsolver_geqrf_ffi` and `hipsolver_orgqr_ffi` from the list of covered targets (as they are added through the test data). ## Test Plan The new back compatibility test data is used in: `tests/export_back_compat_test.py::CompatTest::test_gpu_qr_solver_geqrf`. There are four variants of the test: - `export_back_compat_test.py::CompatTest::test_gpu_qr_solver_geqrf_dtype=f32` - `export_back_compat_test.py::CompatTest::test_gpu_qr_solver_geqrf_dtype=f64` - `export_back_compat_test.py::CompatTest::test_gpu_qr_solver_geqrf_dtype=c64` - `export_back_compat_test.py::CompatTest::test_gpu_qr_solver_geqrf_dtype=c128` ## Test Results All four variants of the `test_gpu_qr_solver_geqrf` unit test pass. Copybara import of the project: -- 153b1feb107d52ff421142bb35635193c6ba3d5b by tsrw2048 <239799652+tsrw2048@users.noreply.github.com>: [ROCm] Added a back compatibility test for GEQRF on ROCm devices. Merging this change closes #34862 COPYBARA_INTEGRATE_REVIEW=https://github.com/jax-ml/jax/pull/34862 from ROCm:rocm-geqrf-compat-test 153b1feb107d52ff421142bb35635193c6ba3d5b PiperOrigin-RevId: 869053499
Author
Parents
Loading