# Source code for gpytorch

#!/usr/bin/env python3

from typing import Optional, Tuple, Union

import linear_operator
import torch
from linear_operator import LinearOperator
from torch import Tensor

from . import (
beta_features,
distributions,
kernels,
lazy,
likelihoods,
means,
metrics,
mlls,
models,
optim,
priors,
settings,
utils,
variational,
)
from .functions import inv_matmul, log_normal_cdf, logdet, matmul  # Deprecated
from .lazy import cat, delazify, lazify
from .mlls import ExactMarginalLogLikelihood
from .module import Module

Anysor = Union[LinearOperator, Tensor]

[docs]def add_diagonal(input: Anysor, diag: Tensor) -> LinearOperator: r""" Adds an element to the diagonal of the matrix :math:\mathbf A. :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param diag: Diagonal to add :return: :math:\mathbf A + \text{diag}(\mathbf d), where :math:\mathbf A is the linear operator and :math:\mathbf d is the diagonal component """ return linear_operator.add_diagonal(input=input, diag=diag)
[docs]def add_jitter(input: Anysor, jitter_val: float = 1e-3) -> Anysor: r""" Adds jitter (i.e., a small diagonal component) to the matrix this LinearOperator represents. This is equivalent to calling :meth:~linear_operator.operators.LinearOperator.add_diagonal with a scalar tensor. :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param jitter_val: The diagonal component to add :return: :math:\mathbf A + \alpha (\mathbf I), where :math:\mathbf A is the linear operator and :math:\alpha is :attr:jitter_val. """ return linear_operator.add_jitter(input=input, jitter_val=jitter_val)
[docs]def diagonalization(input: Anysor, method: Optional[str] = None) -> Tuple[Tensor, Tensor]: r""" Returns a (usually partial) diagonalization of a symmetric positive definite matrix (or batch of matrices). :math:\mathbf A. Options are either "lanczos" or "symeig". "lanczos" runs Lanczos while "symeig" runs LinearOperator.symeig. :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param method: Specify the method to use ("lanczos" or "symeig"). The method will be determined based on size if not specified. :return: eigenvalues and eigenvectors representing the diagonalization. """ return linear_operator.diagonalization(input=input, method=method)
[docs]def dsmm( sparse_mat: Union[torch.sparse.HalfTensor, torch.sparse.FloatTensor, torch.sparse.DoubleTensor], dense_mat: Tensor, ) -> Tensor: r""" Performs the (batch) matrix multiplication :math:\mathbf{SD} where :math:\mathbf S is a sparse matrix and :math:\mathbf D is a dense matrix. :param sparse_mat: Sparse matrix :math:\mathbf S (... x M x N) :param dense_mat: Dense matrix :math:\mathbf D (... x N x O) :return: :math:\mathbf S \mathbf D (... x M x N) """ return linear_operator.dsmm(sparse_mat=sparse_mat, dense_mat=dense_mat)
[docs]def inv_quad(input: Anysor, inv_quad_rhs: Tensor, reduce_inv_quad: bool = True) -> Tensor: r""" Computes an inverse quadratic form (w.r.t self) with several right hand sides, i.e: .. math:: \text{tr}\left( \mathbf R^\top \mathbf A^{-1} \mathbf R \right), where :math:\mathbf A is a positive definite matrix (or batch of matrices) and :math:\mathbf R represents the right hand sides (:attr:inv_quad_rhs). If :attr:reduce_inv_quad is set to false (and :attr:inv_quad_rhs is supplied), the function instead computes .. math:: \text{diag}\left( \mathbf R^\top \mathbf A^{-1} \mathbf R \right). :param input: :math:\mathbf A - the positive definite matrix (... X N X N) :param inv_quad_rhs: :math:\mathbf R - the right hand sides of the inverse quadratic term (... x N x M) :param reduce_inv_quad: Whether to compute :math:\text{tr}\left( \mathbf R^\top \mathbf A^{-1} \mathbf R \right) or :math:\text{diag}\left( \mathbf R^\top \mathbf A^{-1} \mathbf R \right). :returns: The inverse quadratic term. If reduce_inv_quad=True, the inverse quadratic term is of shape (...). Otherwise, it is (... x M). """ return linear_operator.inv_quad(input=input, inv_quad_rhs=inv_quad_rhs, reduce_inv_quad=reduce_inv_quad)
[docs]def inv_quad_logdet( input: Anysor, inv_quad_rhs: Optional[Tensor] = None, logdet: bool = False, reduce_inv_quad: bool = True ) -> Tuple[Tensor, Tensor]: r""" Calls both :func:inv_quad_logdet and :func:logdet on a positive definite matrix (or batch) :math:\mathbf A. However, calling this method is far more efficient and stable than calling each method independently. :param input: :math:\mathbf A - the positive definite matrix (... X N X N) :param inv_quad_rhs: :math:\mathbf R - the right hand sides of the inverse quadratic term (... x N x M) :param logdet: Whether or not to compute the logdet term :math:\log \vert \mathbf A \vert. :param reduce_inv_quad: Whether to compute :math:\text{tr}\left( \mathbf R^\top \mathbf A^{-1} \mathbf R \right) or :math:\text{diag}\left( \mathbf R^\top \mathbf A^{-1} \mathbf R \right). :returns: The inverse quadratic term (or None), and the logdet term (or None). If reduce_inv_quad=True, the inverse quadratic term is of shape (...). Otherwise, it is (... x M). """ return linear_operator.inv_quad_logdet( input=input, inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad )
[docs]def pivoted_cholesky( input: Anysor, rank: int, error_tol: Optional[float] = None, return_pivots: bool = False ) -> Union[Tensor, Tuple[Tensor, Tensor]]: r""" Performs a partial pivoted Cholesky factorization of a positive definite matrix (or batch of matrices). :math:\mathbf L \mathbf L^\top = \mathbf A. The partial pivoted Cholesky factor :math:\mathbf L \in \mathbb R^{N \times \text{rank}} forms a low rank approximation to the LinearOperator. The pivots are selected greedily, correspoading to the maximum diagonal element in the residual after each Cholesky iteration. See Harbrecht et al., 2012_. :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param rank: The size of the partial pivoted Cholesky factor. :param error_tol: Defines an optional stopping criterion. If the residual of the factorization is less than :attr:error_tol, then the factorization will exit early. This will result in a :math:\leq \text{ rank} factor. :param return_pivots: Whether or not to return the pivots alongside the partial pivoted Cholesky factor. :return: The ... x N x rank factor (and optionally the ... x N pivots if :attr:return_pivots is True). .. _Harbrecht et al., 2012: https://www.sciencedirect.com/science/article/pii/S0168927411001814 """ return linear_operator.pivoted_cholesky(input=input, rank=rank, return_pivots=return_pivots)
[docs]def root_decomposition(input: Anysor, method: Optional[str] = None) -> LinearOperator: r""" Returns a (usually low-rank) root decomposition linear operator of the positive definite matrix (or batch of matrices) :math:\mathbf A. This can be used for sampling from a Gaussian distribution, or for obtaining a low-rank version of a matrix. :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param method: Which method to use to perform the root decomposition. Choices are: "cholesky", "lanczos", "symeig", "pivoted_cholesky", or "svd". :return: A tensor :math:\mathbf R such that :math:\mathbf R \mathbf R^\top \approx \mathbf A. """ return linear_operator.root_decomposition(input=input, method=method)
[docs]def root_inv_decomposition( input: Anysor, initial_vectors: Optional[Tensor] = None, test_vectors: Optional[Tensor] = None, method: Optional[str] = None, ) -> LinearOperator: r""" Returns a (usually low-rank) inverse root decomposition linear operator of the PSD LinearOperator :math:\mathbf A. This can be used for sampling from a Gaussian distribution, or for obtaining a low-rank version of a matrix. The root_inv_decomposition is performed using a partial Lanczos tridiagonalization. :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param initial_vectors: Vectors used to initialize the Lanczos decomposition. The best initialization vector (determined by :attr:test_vectors) will be chosen. :param test_vectors: Vectors used to test the accuracy of the decomposition. :param method: Root decomposition method to use (symeig, diagonalization, lanczos, or cholesky). :return: A tensor :math:\mathbf R such that :math:\mathbf R \mathbf R^\top \approx \mathbf A^{-1}. """ return linear_operator.root_inv_decomposition( input=input, initial_vectors=initial_vectors, test_vectors=test_vectors, method=method )
[docs]def solve(input: Anysor, rhs: Tensor, lhs: Optional[Tensor] = None) -> Tensor: r""" Given a positive definite matrix (or batch of matrices) :math:\mathbf A, computes a linear solve with right hand side :math:\mathbf R: .. math:: \mathbf A^{-1} \mathbf R, where :math:\mathbf R is :attr:right_tensor and :math:\mathbf A is the LinearOperator. .. note:: Unlike :func:torch.linalg.solve, this function can take an optional :attr:left_tensor attribute. If this is supplied :func:gpytorch.solve computes .. math:: \mathbf L \mathbf A^{-1} \mathbf R, where :math:\mathbf L is :attr:left_tensor. Supplying this can reduce the number of solver calls required in the backward pass. :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param rhs: :math:\mathbf R - the right hand side :param lhs: :math:\mathbf L - the left hand side :return: :math:\mathbf A^{-1} \mathbf R or :math:\mathbf L \mathbf A^{-1} \mathbf R. """ return linear_operator.solve(input=input, rhs=rhs, lhs=lhs)
[docs]def sqrt_inv_matmul(input: Anysor, rhs: Tensor, lhs: Optional[Tensor] = None) -> Tensor: r""" Given a positive definite matrix (or batch of matrices) :math:\mathbf A and a right hand size :math:\mathbf R, computes .. math:: \mathbf A^{-1/2} \mathbf R, If :attr:lhs is supplied, computes .. math:: \mathbf L \mathbf A^{-1/2} \mathbf R, where :math:\mathbf L is :attr:lhs. (Supplying :attr:lhs can reduce the number of solver calls required in the backward pass.) :param input: The matrix (or batch of matrices) :math:\mathbf A (... x N x N). :param rhs: :math:\mathbf R - the right hand side :param lhs: :math:\mathbf L - the left hand side :return: :math:\mathbf A^{-1/2} \mathbf R or :math:\mathbf L \mathbf A^{-1/2} \mathbf R. """ return linear_operator.sqrt_inv_matmul(input=input, rhs=rhs, lhs=lhs)
# Read version number as written by setuptools_scm try: from gpytorch.version import version as __version__ except Exception: # pragma: no cover __version__ = "Unknown" # pragma: no cover __all__ = [ # Submodules "distributions", "kernels", "lazy", "likelihoods", "means", "metrics", "mlls", "models", "optim", "priors", "utils", "variational", # Classes "Module", "ExactMarginalLogLikelihood", # Functions "add_diagonal", "add_jitter", "dsmm", "inv_quad", "inv_quad_logdet", "pivoted_cholesky", "root_decomposition", "root_inv_decomposition", "solve", "sqrt_inv_matmul", # Context managers "beta_features", "settings", # Other "__version__", # Deprecated "add_diag", "cat", "delazify", "inv_matmul", "lazify", "logdet", "log_normal_cdf", "matmul", ]