torch.linalg.matrix_sqrth

torch.linalg.matrix_sqrth(A) → Tensor

Computes the principal square root of a symmetric (resp. Hermitian) positive-definite matrix.

Letting $\mathbb{K}$ be $\mathbb{R}$ or $\mathbb{C}$, for a symmetric (resp. Hermitian) positive-definite matrix $A \in \mathbb{K}^{n \times n}$, this function returns the unique symmetric (resp. Hermitian) positive-definite matrix $X \in \mathbb{K}^{n \times n}$ such that

$$XX = A.$$

Supports input of float, double, cfloat and cdouble dtypes. Also supports batches of matrices, and if A is a batch of matrices then the output has the same batch dimensions.

Note

Only the lower triangular part of A is used in the computation, and A is assumed to be symmetric (resp. Hermitian). See torch.linalg.eigh().

See also

torch.linalg.cholesky() computes a different factorization of a symmetric (resp. Hermitian) positive-definite matrix.

Parameters:

A (Tensor) – tensor of shape (*, n, n) where * is zero or more batch dimensions consisting of symmetric (resp. Hermitian) positive-definite matrices.

Examples:

>>> A = torch.tensor([[2., 0.], [0., 9.]])
>>> torch.linalg.matrix_sqrth(A)
tensor([[1.4142, 0.0000],
        [0.0000, 3.0000]])

>>> A = torch.randn(2, 3, 3)
>>> A = A @ A.mT + 3 * torch.eye(3)  # batch of symmetric positive-definite matrices
>>> X = torch.linalg.matrix_sqrth(A)
>>> torch.allclose(X @ X, A, atol=1e-5)
True