- ott.math.matrix_square_root.sqrtm(x, threshold=1e-06, min_iterations=0, inner_iterations=10, max_iterations=1000, regularization=1e-06) = <jax._src.custom_derivatives.custom_vjp object>[source]#
Higham algorithm to compute matrix square root of p.d. matrix.
See [Higham, 1997], eq. 2.6b
Array) – a (batch of) square p.s.d. matrices of the same size.
float) – convergence tolerance threshold for Newton-Schulz iterations.
int) – min number of iterations after which error is computed.
int) – error is re-evaluated every inner_iterations iterations.
int) – max number of iterations.
float) – small regularizer added to norm of x, before normalization.
- Return type:
Square root matrix of x (or x’s if batch), its inverse, errors along iterates.