tttc3 c89be05b5b Enable pivoted QR on CPU devices. ...

A pivoted QR factorization is possible in `scipy.linalg.qr`, thanks
to the `geqp3` routine of LAPACK. To provide the same functionality
in JAX, we implement a new primitive `geqp3_p` which calls the LAPACK
routine via the FFI on CPU devices.

Both `jax.scipy.linalg.qr` and `jax.lax.linalg.qr` now support the
use of column-pivoting on CPU devices.

To provide a GPU implementation of `geqp3` may require using MAGMA,
due to the lack of a `geqp3` implementation in `cuSolver` -  see
ccb331707e80b16d89de6e5c9f2f89b87c1682ed (`jax.lax.linalg.eig`) for
an example of using MAGMA in GPU lowerings. Such a GPU implementation
can be considered in the future.

2025-01-09 20:44:45 +00: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

Enable pivoted QR on CPU devices.

2025-01-09 20:44:45 +00:00

lapack_kernels_using_lapack.cc

Enable pivoted QR on CPU devices.

2025-01-09 20:44:45 +00:00

lapack_kernels.cc

Enable pivoted QR on CPU devices.

2025-01-09 20:44:45 +00:00

lapack_kernels.h

Enable pivoted QR on CPU devices.

2025-01-09 20:44:45 +00:00

lapack.cc

Enable pivoted QR on CPU devices.

2025-01-09 20:44:45 +00:00