quantax.optimizer.lstsq_shift_eig#
- quantax.optimizer.lstsq_shift_eig(rshift: float | None = None, ashift: float = 1e-06, dtype: str | type[Any] | dtype | SupportsDType | None = None, *, jaxmg_ndevices: int = 1) Callable[[...], Array]#
The SR branch of
auto_shift_eig, computing \(x = (A^†A + \epsilon I)^{-1} A^†b\) by directly forming the shifted matrix and solving it with a Cholesky solver. Suitable for overdetermined problems where the number of samples exceeds the number of parameters.- Parameters:
rshift – The relative diagonal shift, entering the shift as \(\epsilon = \mathrm{Tr}(A^†A) \times \mathrm{rshift} / \sqrt{m} + \mathrm{ashift}\). Default to \(10^{-12}\) for double precision and \(10^{-6}\) for single precision.
ashift – The absolute diagonal shift, default to 1e-6.
dtype – The dtype used internally in the solver. By default, real-valued inputs use float64 and complex-valued inputs use complex128.
jaxmg_ndevices – The number of devices to use with jaxmg for distributed linear algebra. By default it is set to 1, which means not using
jaxmg. Setting it to the number of devices per node will enablejaxmg, which is often used for large-scale problems where the matrix is too large to fit in memory on a single device.
- Returns:
A solver function with two arguments A and b and one output x as the solution of \(A x = b\).