jax.scipy.linalg.sqrtm#

jax.scipy.linalg.sqrtm(A, blocksize=1)[source]#

Matrix square root.

LAX-backend implementation of scipy.linalg._matfuncs_sqrtm.sqrtm().

This differs from scipy.linalg.sqrtm in that the return type of jax.scipy.linalg.sqrtm is always complex64 for 32-bit input, and complex128 for 64-bit input.

This function implements the complex Schur method described in [A]. It does not use recursive blocking to speed up computations as a Sylvester Equation solver is not available yet in JAX.

[A] Björck, Å., & Hammarling, S. (1983).

“A Schur method for the square root of a matrix”. Linear algebra and its applications, 52, 127-140.

Original docstring below.

Parameters

  • A ((N, N) array_like) – Matrix whose square root to evaluate

  • blocksize (integer, optional) – If the blocksize is not degenerate with respect to the size of the input array, then use a blocked algorithm. (Default: 64)

Return type

Array

Returns

  • sqrtm ((N, N) ndarray) – Value of the sqrt function at A

  • errest (float) – (if disp == False)

    Frobenius norm of the estimated error, ||err||_F / ||A||_F

References

1

Edvin Deadman, Nicholas J. Higham, Rui Ralha (2013) “Blocked Schur Algorithms for Computing the Matrix Square Root, Lecture Notes in Computer Science, 7782. pp. 171-182.