ott.math.matrix_square_root.sqrtm

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>

Higham algorithm to compute matrix square root of p.d. matrix.

See [Higham, 1997], eq. 2.6b

Parameters:

• x (Array) – a (batch of) square p.s.d. matrices of the same size.

• threshold (float) – convergence tolerance threshold for Newton-Schulz iterations.

• min_iterations (int) – min number of iterations after which error is computed.

• inner_iterations (int) – error is re-evaluated every inner_iterations iterations.

• max_iterations (int) – max number of iterations.

• regularization (float) – small regularizer added to norm of x, before normalization.

Return type:

Tuple[Array, Array, Array]

Returns:

Square root matrix of x (or x’s if batch), its inverse, errors along iterates.