Source code for tensorplay.linalg._decompositions

"""Matrix factorizations.

The kernels behind these bind to the CPU and CUDA decomposition backends;
the wrappers add the named result tuples and the mode handling.
"""
import tensorplay
from tensorplay import _C
from tensorplay._C import (
    linalg_cholesky,
    linalg_cholesky_ex,
    linalg_eig,
    linalg_eigh,
    linalg_eigvals,
    linalg_eigvalsh,
    linalg_householder_product as householder_product,
    linalg_ldl_factor as ldl_factor,
    linalg_ldl_factor_ex as ldl_factor_ex,
    linalg_lu as lu,
    linalg_lu_factor as lu_factor,
    linalg_lu_factor_ex as lu_factor_ex,
)

from ._common import CholeskyExResult, EigResult, EighResult, QRResult, SVDResult, check_floating

__all__ = [
    "cholesky",
    "cholesky_ex",
    "eig",
    "eigh",
    "eigvals",
    "eigvalsh",
    "householder_product",
    "ldl_factor",
    "ldl_factor_ex",
    "lu",
    "lu_factor",
    "lu_factor_ex",
    "polar",
    "qr",
    "svd",
    "svdvals",
]



[docs]
def cholesky(A, *, upper=False):
    """cholesky(A, *, upper=False) -> Tensor"""
    return linalg_cholesky(A, upper=upper)




[docs]
def cholesky_ex(A, *, upper=False, check_errors=False):
    """cholesky_ex(A, *, upper=False, check_errors=False) -> CholeskyExResult(L, info)"""
    L, info = linalg_cholesky_ex(A, upper=upper, check_errors=check_errors)
    return CholeskyExResult(L, info)




[docs]
def eigh(A, UPLO="L"):
    """eigh(A, UPLO='L') -> EighResult(eigenvalues, eigenvectors)"""
    UPLO = str(UPLO).upper()
    if UPLO not in ("L", "U"):
        raise ValueError("linalg.eigh: UPLO must be 'L' or 'U'")
    values, vectors = linalg_eigh(A, UPLO)
    return EighResult(values, vectors)




[docs]
def eigvalsh(A, UPLO="L"):
    """eigvalsh(A, UPLO='L') -> Tensor"""
    UPLO = str(UPLO).upper()
    if UPLO not in ("L", "U"):
        raise ValueError("linalg.eigvalsh: UPLO must be 'L' or 'U'")
    return linalg_eigvalsh(A, UPLO)




[docs]
def eig(A):
    """eig(A) -> EigResult(eigenvalues, eigenvectors)"""
    values, vectors = linalg_eig(A)
    return EigResult(values, vectors)




[docs]
def eigvals(A):
    """eigvals(A) -> Tensor"""
    return linalg_eigvals(A)




[docs]
def svd(A, full_matrices=True, *, driver=None):
    """svd(A, full_matrices=True, *, driver=None) -> SVDResult(U, S, Vh)"""
    U, S, Vh = _C.linalg_svd(A, full_matrices, driver=driver)
    return SVDResult(U, S, Vh)




[docs]
def svdvals(A, *, driver=None):
    """svdvals(A, *, driver=None) -> Tensor"""
    return _C.linalg_svdvals(A, driver=driver)




[docs]
def qr(A, mode="reduced"):
    """qr(A, mode='reduced') -> QRResult(Q, R)"""
    if mode == "R":
        mode = "r"
    if mode not in ("reduced", "complete", "r"):
        raise ValueError(
            "linalg.qr: mode must be 'reduced', 'complete', or 'r'")
    Q, R = _C.linalg_qr(A, mode)
    if mode in ("r", "R"):
        empty = tensorplay.empty(
            list(A.shape[:-2]) + [A.shape[-2], 0],
            dtype=A.dtype,
            device=A.device,
        )
        return QRResult(empty, R)
    return QRResult(Q, R)




[docs]
def polar(A):
    """polar(A) -> (Tensor Q, Tensor R) with A = Q R"""
    check_floating(A, "polar")
    if A.dim() < 2 or A.shape[-2] < A.shape[-1]:
        raise ValueError(
            "linalg.polar: input must have at least as many rows as columns")
    U, S, Vh = svd(A, full_matrices=False)
    Q = U @ Vh
    V = _C.conj_physical(Vh).transpose(-2, -1)
    R = V @ (S.unsqueeze(-1) * Vh)
    return Q, R