Shortcuts

Source code for torch._lowrank

"""Implement various linear algebra algorithms for low rank matrices."""

__all__ = ["svd_lowrank", "pca_lowrank"]

from typing import Optional, Tuple

import torch
from torch import _linalg_utils as _utils, Tensor
from torch.overrides import handle_torch_function, has_torch_function


def get_approximate_basis(
    A: Tensor,
    q: int,
    niter: Optional[int] = 2,
    M: Optional[Tensor] = None,
) -> Tensor:
    """Return tensor :math:`Q` with :math:`q` orthonormal columns such
    that :math:`Q Q^H A` approximates :math:`A`. If :math:`M` is
    specified, then :math:`Q` is such that :math:`Q Q^H (A - M)`
    approximates :math:`A - M`. without instantiating any tensors
    of the size of :math:`A` or :math:`M`.

    .. note:: The implementation is based on the Algorithm 4.4 from
              Halko et al., 2009.

    .. note:: For an adequate approximation of a k-rank matrix
              :math:`A`, where k is not known in advance but could be
              estimated, the number of :math:`Q` columns, q, can be
              choosen according to the following criteria: in general,
              :math:`k <= q <= min(2*k, m, n)`. For large low-rank
              matrices, take :math:`q = k + 5..10`.  If k is
              relatively small compared to :math:`min(m, n)`, choosing
              :math:`q = k + 0..2` may be sufficient.

    .. note:: To obtain repeatable results, reset the seed for the
              pseudorandom number generator

    Args::
        A (Tensor): the input tensor of size :math:`(*, m, n)`

        q (int): the dimension of subspace spanned by :math:`Q`
                 columns.

        niter (int, optional): the number of subspace iterations to
                               conduct; ``niter`` must be a
                               nonnegative integer. In most cases, the
                               default value 2 is more than enough.

        M (Tensor, optional): the input tensor's mean of size
                              :math:`(*, m, n)`.

    References::
        - Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding
          structure with randomness: probabilistic algorithms for
          constructing approximate matrix decompositions,
          arXiv:0909.4061 [math.NA; math.PR], 2009 (available at
          `arXiv <http://arxiv.org/abs/0909.4061>`_).
    """

    niter = 2 if niter is None else niter
    dtype = _utils.get_floating_dtype(A) if not A.is_complex() else A.dtype
    matmul = _utils.matmul

    R = torch.randn(A.shape[-1], q, dtype=dtype, device=A.device)

    # The following code could be made faster using torch.geqrf + torch.ormqr
    # but geqrf is not differentiable

    X = matmul(A, R)
    if M is not None:
        X = X - matmul(M, R)
    Q = torch.linalg.qr(X).Q
    for i in range(niter):
        X = matmul(A.mH, Q)
        if M is not None:
            X = X - matmul(M.mH, Q)
        Q = torch.linalg.qr(X).Q
        X = matmul(A, Q)
        if M is not None:
            X = X - matmul(M, Q)
        Q = torch.linalg.qr(X).Q
    return Q


[docs]def svd_lowrank( A: Tensor, q: Optional[int] = 6, niter: Optional[int] = 2, M: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Tensor]: r"""Return the singular value decomposition ``(U, S, V)`` of a matrix, batches of matrices, or a sparse matrix :math:`A` such that :math:`A \approx U \operatorname{diag}(S) V^{\text{H}}`. In case :math:`M` is given, then SVD is computed for the matrix :math:`A - M`. .. note:: The implementation is based on the Algorithm 5.1 from Halko et al., 2009. .. note:: For an adequate approximation of a k-rank matrix :math:`A`, where k is not known in advance but could be estimated, the number of :math:`Q` columns, q, can be choosen according to the following criteria: in general, :math:`k <= q <= min(2*k, m, n)`. For large low-rank matrices, take :math:`q = k + 5..10`. If k is relatively small compared to :math:`min(m, n)`, choosing :math:`q = k + 0..2` may be sufficient. .. note:: This is a randomized method. To obtain repeatable results, set the seed for the pseudorandom number generator .. note:: In general, use the full-rank SVD implementation :func:`torch.linalg.svd` for dense matrices due to its 10x higher performance characteristics. The low-rank SVD will be useful for huge sparse matrices that :func:`torch.linalg.svd` cannot handle. Args:: A (Tensor): the input tensor of size :math:`(*, m, n)` q (int, optional): a slightly overestimated rank of A. niter (int, optional): the number of subspace iterations to conduct; niter must be a nonnegative integer, and defaults to 2 M (Tensor, optional): the input tensor's mean of size :math:`(*, m, n)`, which will be broadcasted to the size of A in this function. References:: - Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding structure with randomness: probabilistic algorithms for constructing approximate matrix decompositions, arXiv:0909.4061 [math.NA; math.PR], 2009 (available at `arXiv <https://arxiv.org/abs/0909.4061>`_). """ if not torch.jit.is_scripting(): tensor_ops = (A, M) if not set(map(type, tensor_ops)).issubset( (torch.Tensor, type(None)) ) and has_torch_function(tensor_ops): return handle_torch_function( svd_lowrank, tensor_ops, A, q=q, niter=niter, M=M ) return _svd_lowrank(A, q=q, niter=niter, M=M)
def _svd_lowrank( A: Tensor, q: Optional[int] = 6, niter: Optional[int] = 2, M: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor, Tensor]: # Algorithm 5.1 in Halko et al., 2009 q = 6 if q is None else q m, n = A.shape[-2:] matmul = _utils.matmul if M is not None: M = M.broadcast_to(A.size()) # Assume that A is tall if m < n: A = A.mH if M is not None: M = M.mH Q = get_approximate_basis(A, q, niter=niter, M=M) B = matmul(Q.mH, A) if M is not None: B = B - matmul(Q.mH, M) U, S, Vh = torch.linalg.svd(B, full_matrices=False) V = Vh.mH U = Q.matmul(U) if m < n: U, V = V, U return U, S, V
[docs]def pca_lowrank( A: Tensor, q: Optional[int] = None, center: bool = True, niter: int = 2, ) -> Tuple[Tensor, Tensor, Tensor]: r"""Performs linear Principal Component Analysis (PCA) on a low-rank matrix, batches of such matrices, or sparse matrix. This function returns a namedtuple ``(U, S, V)`` which is the nearly optimal approximation of a singular value decomposition of a centered matrix :math:`A` such that :math:`A \approx U \operatorname{diag}(S) V^{\text{H}}` .. note:: The relation of ``(U, S, V)`` to PCA is as follows: - :math:`A` is a data matrix with ``m`` samples and ``n`` features - the :math:`V` columns represent the principal directions - :math:`S ** 2 / (m - 1)` contains the eigenvalues of :math:`A^T A / (m - 1)` which is the covariance of ``A`` when ``center=True`` is provided. - ``matmul(A, V[:, :k])`` projects data to the first k principal components .. note:: Different from the standard SVD, the size of returned matrices depend on the specified rank and q values as follows: - :math:`U` is m x q matrix - :math:`S` is q-vector - :math:`V` is n x q matrix .. note:: To obtain repeatable results, reset the seed for the pseudorandom number generator Args: A (Tensor): the input tensor of size :math:`(*, m, n)` q (int, optional): a slightly overestimated rank of :math:`A`. By default, ``q = min(6, m, n)``. center (bool, optional): if True, center the input tensor, otherwise, assume that the input is centered. niter (int, optional): the number of subspace iterations to conduct; niter must be a nonnegative integer, and defaults to 2. References:: - Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp, Finding structure with randomness: probabilistic algorithms for constructing approximate matrix decompositions, arXiv:0909.4061 [math.NA; math.PR], 2009 (available at `arXiv <http://arxiv.org/abs/0909.4061>`_). """ if not torch.jit.is_scripting(): if type(A) is not torch.Tensor and has_torch_function((A,)): return handle_torch_function( pca_lowrank, (A,), A, q=q, center=center, niter=niter ) (m, n) = A.shape[-2:] if q is None: q = min(6, m, n) elif not (q >= 0 and q <= min(m, n)): raise ValueError( f"q(={q}) must be non-negative integer and not greater than min(m, n)={min(m, n)}" ) if not (niter >= 0): raise ValueError(f"niter(={niter}) must be non-negative integer") dtype = _utils.get_floating_dtype(A) if not center: return _svd_lowrank(A, q, niter=niter, M=None) if _utils.is_sparse(A): if len(A.shape) != 2: raise ValueError("pca_lowrank input is expected to be 2-dimensional tensor") c = torch.sparse.sum(A, dim=(-2,)) / m # reshape c column_indices = c.indices()[0] indices = torch.zeros( 2, len(column_indices), dtype=column_indices.dtype, device=column_indices.device, ) indices[0] = column_indices C_t = torch.sparse_coo_tensor( indices, c.values(), (n, 1), dtype=dtype, device=A.device ) ones_m1_t = torch.ones(A.shape[:-2] + (1, m), dtype=dtype, device=A.device) M = torch.sparse.mm(C_t, ones_m1_t).mT return _svd_lowrank(A, q, niter=niter, M=M) else: C = A.mean(dim=(-2,), keepdim=True) return _svd_lowrank(A - C, q, niter=niter, M=None)

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources