Activate QR Factorization to XLA's FFI

PiperOrigin-RevId: 666722604
jax-ml · Aug 23, 2024 · c430b0c · c430b0c
1 parent 07767e8
commit c430b0c
Show file tree

Hide file tree

Showing 5 changed files with 556 additions and 38 deletions.
diff --git a/jax/_src/export/_export.py b/jax/_src/export/_export.py
@@ -923,9 +923,22 @@ def _check_lowering(lowering) -> None:
  "\n".join(not_implemented_msgs))
 
 _CPU_FFI_KERNELS = [
- "lapack_spotrf_ffi", "lapack_dpotrf_ffi", "lapack_cpotrf_ffi", "lapack_zpotrf_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_spotrf_ffi",
+ "lapack_dpotrf_ffi",
+ "lapack_cpotrf_ffi",
+ "lapack_zpotrf_ffi",
+ "lapack_sgeqrf_ffi",
+ "lapack_dgeqrf_ffi",
+ "lapack_cgeqrf_ffi",
+ "lapack_zgeqrf_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",
 ]
 # These are the JAX custom call target names that are guaranteed to be stable.
 # Their backwards compatibility is tested by back_compat_test.py.