Activate Tridiagonal Reduction to XLA's FFI

Additionally, created a missing backward compatibility test for the old LAPACK kernels of Tridiagonal Reduction. PiperOrigin-RevId: 704234350
jax-ml · Dec 9, 2024 · d474fed · d474fed
1 parent adb2bf6
commit d474fed
Show file tree

Hide file tree

Showing 5 changed files with 975 additions and 48 deletions.
diff --git a/jax/_src/export/_export.py b/jax/_src/export/_export.py
@@ -997,6 +997,7 @@ def _check_lowering(lowering) -> None:
     "lapack_sgeev_ffi", "lapack_dgeev_ffi", "lapack_cgeev_ffi", "lapack_zgeev_ffi",
     "lapack_sgesdd_ffi", "lapack_dgesdd_ffi", "lapack_cgesdd_ffi", "lapack_zgesdd_ffi",
     "lapack_sgetrf_ffi", "lapack_dgetrf_ffi", "lapack_cgetrf_ffi", "lapack_zgetrf_ffi",
+    "lapack_ssytrd_ffi", "lapack_dsytrd_ffi", "lapack_chetrd_ffi", "lapack_zhetrd_ffi",
     "lapack_sgehrd_ffi", "lapack_dgehrd_ffi", "lapack_cgehrd_ffi", "lapack_zgehrd_ffi",
     "lapack_sgees_ffi", "lapack_dgees_ffi", "lapack_cgees_ffi", "lapack_zgees_ffi",
 ]
@@ -1022,6 +1023,8 @@ def _check_lowering(lowering) -> None:
     "blas_strsm", "blas_dtrsm", "blas_ctrsm", "blas_ztrsm",
     # schur on CPU
     "lapack_sgees", "lapack_dgees", "lapack_cgees", "lapack_zgees",
+    # tridiagonal on CPU
+    "lapack_ssytrd", "lapack_dsytrd", "lapack_chetrd", "lapack_zhetrd",
     # hessenberg on CPU
     "lapack_sgehrd", "lapack_dgehrd", "lapack_cgehrd", "lapack_zgehrd",
     # lu on GPU