mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jaxlib/cpu
History
Dan Foreman-Mackey c1de7c733d Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU. 
In implementing https://github.com/jax-ml/jax/pull/25787, I realized that while we lower `tridiagonal_solve` to cuSPARSE on GPU, we were using an explicit implementation of the Thomas algorithm on CPU. We should instead lower to LAPACK's `gtsv` on CPU because it should be more numerically stable and faster.

PiperOrigin-RevId: 714069225
2025-01-10 08:56:46 -08:00
..
_lapack
Port Schur Decomposition to XLA's FFI
2024-10-14 06:46:42 -07:00
BUILD
Rename pybind_extension to nanobind_extension.
2025-01-06 11:53:44 -08:00
cpu_kernels.cc
Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU.
2025-01-10 08:56:46 -08:00
lapack_kernels_using_lapack.cc
Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU.
2025-01-10 08:56:46 -08:00
lapack_kernels.cc
Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU.
2025-01-10 08:56:46 -08:00
lapack_kernels.h
Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU.
2025-01-10 08:56:46 -08:00
lapack.cc
Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU.
2025-01-10 08:56:46 -08:00