quantax.optimizer.auto_shift_eig#
- quantax.optimizer.auto_shift_eig(rshift: float | None = None, ashift: float = 1e-06, dtype: str | type[Any] | dtype | SupportsDType | None = None, *, jaxmg_ndevices: int = 1) Callable[[...], Array]#
Obtain the least-square minimum-norm solver for the linear equation \(Ax=b\) using diagonal shift. It automatically chooses between \(x = (A^†A)^{-1} A^†b\) and \(x = A^†(A A^†)^{-1} b\), which respectively correspond to SR and MinSR.
Given \(M = A^†A\) or \(M = A A^†\) of dimension \(n\), the diagonal shift modifies it to \(M' = M + \epsilon I\) for stable inversion, with \(\epsilon = \mathrm{Tr}(M) \times \mathrm{rshift} / \sqrt{n} + \mathrm{ashift}\).
- Parameters:
rshift – The relative diagonal shift. Default to be \(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. This option is often used for large-scale problems where the matrix is too large to fit in memory on a single device. It requiresjaxmgto be installed and properly configured.
- Returns:
A solver function with two arguments A and b and one output x as the solution of \(A x = b\).