Package: linear-operator
LinearOperatorLinearOperator is a PyTorch package for abstracting away the linear algebra routines needed for structured matrices (or operators).This package is in beta.Currently, most of the functionality only supports positive semi-definite and triangular matrices.
Package development TODOs:Support PSD operatorsSupport triangular operatorsInterface to specify structure (i.e. symmetric, triangular, PSD, etc.)Add algebraic routines for symmetric operatorsAdd algebraic routines for generic square operatorsAdd algebraic routines for generic rectangular operatorsAdd sparse operatorsTo get started, run eitherpipinstalllinear_operator# orcondainstalllinear_operator-cgpytorchorsee belowfor more detailed instructions.Why LinearOperatorBefore describing what linear operators are and why they make a useful abstraction, it's easiest to see an example.
Let's say you wanted to compute a matrix solve:$$\boldsymbol A^{-1} \boldsymbol b.$$If you didn't know anything about the matrix $\boldsymbol A$, the simplest (and best) way to accomplish this in code is:# A = torch.randn(1000, 1000)# b = torch.randn(1000)torch.linalg.solve(A,b)# computes A^{-1} bWhile this is easy, thesolveroutine is $\mathcal O(N^3)$, which gets very slow as $N$ grows large.However, let's imagine that we knew that $\boldsymbol A$ was equal to a low rank matrix plus a diagonal
(i.e. $\boldsymbol A = \boldsymbol C \boldsymbol C^\top + \boldsymbol D$
for some skinny matrix $\boldsymbol C$ and some diagonal matrix $\boldsymbol D$.)
There's now a very efficient $\boldsymbol O(N)$ routine to compute $\boldsymbol A^{-1}$ (theWoodbury formula).In general, if we know that $\boldsymbol A$ has structure,
we want to use efficient linear algebra routines - rather than the general routines -
that exploit this structure.Without LinearOperatorImplementing the efficient solve that exploits $\boldsymbol A$'s low-rank-plus-diagonal structure would look something like this:deflow_rank_plus_diagonal_solve(C,d,b):# A = C C^T + diag(d)# A^{-1} b = D^{-1} b - D^{-1} C (I + C^T D^{-1} C)^{-1} C^T D^{-1} b#   where D = diag(d)D_inv_b=b/dD_inv_C=C/d.unsqueeze(-1)eye=torch.eye(C.size(-2))return(D_inv_b-D_inv_C@torch.cholesky_solve(C.mT@D_inv_b,torch.linalg.cholesky(eye+C.mT@D_inv_C,upper=False),upper=False))# C = torch.randn(1000, 20)# d = torch.randn(1000)# b = torch.randn(1000)low_rank_plus_diagonal_solve(C,d,b)# computes A^{-1} b in O(N) time, instead of O(N^3)While this is efficient code, it's not ideal for a number of reasons:It's a lot more complicated thantorch.linalg.solve(A, b).There's no object that represents $\boldsymbol A$.
To perform any math with $\boldsymbol A$, we have to pass around the matrixCand the vectord.With LinearOperatorThe LinearOperator package offers the best of both worlds:fromlinear_operator.operatorsimportDiagLinearOperator,LowRankRootLinearOperator# C = torch.randn(1000, 20)# d = torch.randn(1000)# b = torch.randn(1000)A=LowRankRootLinearOperator(C)+DiagLinearOperator(d)# represents C C^T + diag(d)it provides an interface that lets us treat $\boldsymbol A$ as if it were a generic tensor,
using the standard PyTorch API:torch.linalg.solve(A,b)# computes A^{-1} b efficiently!Under-the-hood, theLinearOperatorobject keeps track of the algebraic structure of $\boldsymbol A$ (low rank plus diagonal)
and determines the most efficient routine to use (the Woodbury formula).
This way, we can get a efficient $\mathcal O(N)$ solve while abstracting away all of the details.Crucially, $\boldsymbol A$ is never explicitly instantiated as a matrix, which makes it possible to scale
to very large operators without running out of memory:# C = torch.randn(10000000, 20)# d = torch.randn(10000000)# b = torch.randn(10000000)A=LowRankRootLinearOperator(C)+DiagLinearOperator(d)# represents a 10M x 10M matrix!torch.linalg.solve(A,b)# computes A^{-1} b efficiently!What is a Linear Operator?A linear operator is a generalization of a matrix.
It is a linear function that is defined in by its application to a vector.
The most common linear operators are (potentially structured) matrices,
where the function applying them to a vector are (potentially efficient)
matrix-vector multiplication routines.In code, aLinearOperatoris a class thatspecifies the tensor(s) needed to define the LinearOperator,specifies a_matmulfunction (how the LinearOperator is applied to a vector),specifies a_sizefunction (how big is the LinearOperator if it is represented as a matrix, or batch of matrices), andspecifies a_transpose_nonbatchfunction (the adjoint of the LinearOperator).(optionally) defines other functions (e.g.logdet,eigh, etc.) to accelerate computations for which efficient sturcture-exploiting routines exist.For example:classDiagLinearOperator(linear_operator.LinearOperator):r"""A LinearOperator representing a diagonal matrix."""def__init__(self,diag):# diag: the vector that defines the diagonal of the matrixself.diag=diagdef_matmul(self,v):returnself.diag.unsqueeze(-1)*vdef_size(self):returntorch.Size([*self.diag.shape,self.diag.size(-1)])def_transpose_nonbatch(self):returnself# Diagonal matrices are symmetric# this function is optional, but it will accelerate computationdeflogdet(self):returnself.diag.log().sum(dim=-1)# ...D=DiagLinearOperator(torch.tensor([1.,2.,3.])# Represents the matrix#   [[1., 0., 0.],#    [0., 2., 0.],#    [0., 0., 3.]]torch.matmul(D,torch.tensor([4.,5.,6.])# Returns [4., 10., 18.]While_matmul,_size, and_transpose_nonbatchmight seem like a limited set of functions,
it turns out that most functions on thetorchandtorch.linalgnamespaces can be efficiently implemented
using only these three primitative functions.Moreover, because_matmulis a linear function, it is very easy to compose linear operators in various ways.
For example: adding two linear operators (SumLinearOperator) just requires adding the output of their_matmulfunctions.
This makes it possible to define very complex compositional structures that still yield efficient linear algebraic routines.Finally,LinearOperatorobjects can be composed with one another, yielding newLinearOperatorobjects and automatically keeping track of algebraic structure after each computation.
As a result, users never need to reason about what efficient linear algebra routines to use  (so long as the input elements defined by the user encode known input structure).See theusing LinearOperator objectssection for more details.Use CasesThere are several use cases for the LinearOperator package.
Here we highlight two general themes:Modular Code for Structured MatricesFor example, let's say that you have a generative model that involves
sampling from a high-dimensional multivariate Gaussian.
This sampling operation will require storing and manipulating a large covariance matrix,
so to speed things up you might want to experiment with different structured
approximations of that covariance matrix.
This is easy with the LinearOperator package.fromgpytorch.distributionsimportMultivariateNormal# variance = torch.randn(10000)cov=DiagLinearOperator(variance)# or# cov = LowRankRootLinearOperator(...) + DiagLinearOperator(...)# or# cov = KroneckerProductLinearOperator(...)# or# cov = ToeplitzLinearOperator(...)# or# ...mvn=MultivariateNormal(torch.zeros(cov.size(-1),cov)# 10000-dimensional MVNmvn.rsample()# returns a 10000-dimensional vectorEfficient Routines for Complex OperatorsMany of the efficient linear algebra routines in LinearOperator are iterative algorithms
based on matrix-vector multiplication.
Since matrix-vector multiplication obeys many nice compositional properties
it is possible to obtain efficient routines for extremely complex compositional LienarOperators:fromlinear_operator.operatorsimportKroneckerProductLinearOperator,RootLinearOperator,ToeplitzLinearOperator# mat1 = 200 x 200 PSD matrix# mat2 = 100 x 100 PSD matrix# vec3 = 20000 vectorA=KroneckerProductLinearOperator(mat1,mat2)+RootLinearOperator(ToeplitzLinearOperator(vec3))# represents a 20000 x 20000 matrixtorch.linalg.solve(A,torch.randn(20000))# Sub O(N^3) routine!Using LinearOperator ObjectsLinearOperator objects share (mostly) the same API astorch.Tensorobjects.
Under the hood, these objects use__torch_function__to dispatch all efficient linear algebra operations
to thetorchandtorch.linalgnamespaces.
This includestorch.addtorch.cattorch.clonetorch.diagonaltorch.dimtorch.divtorch.expandtorch.logdettorch.matmultorch.numeltorch.permutetorch.prodtorch.squeezetorch.subtorch.sumtorch.transposetorch.unsqueezetorch.linalg.choleskytorch.linalg.eightorch.linalg.eigvalshtorch.linalg.solvetorch.linalg.svdEach of these functions will either return atorch.Tensor, or a newLinearOperatorobject,
depending on the function.
For example:# A = RootLinearOperator(...)# B = ToeplitzLinearOperator(...)# d = vecC=torch.matmul(A,B)# A new LienearOperator representing the product of A and Btorch.linalg.solve(C,d)# A torch.TensorFor more examples, see theexamples folder.Batch Support and BroadcastingLinearOperatorobjects operate naturally in batch mode.
For example, to represent a batch of 3100 x 100diagonal matrices:# d = torch.randn(3, 100)D=DiagLinearOperator(d)# Reprents an operator of size 3 x 100 x 100These objects fully support broadcasted operations:D@torch.randn(100,2)# Returns a tensor of size 3 x 100 x 2D2=DiagLinearOperator(torch.randn([2,1,100]))# Represents an operator of size 2 x 1 x 100 x 100D2+D# Represents an operator of size 2 x 3 x 100 x 100IndexingLinearOperatorobjects can be indexed in ways similar to torch Tensors. This includes:Integer indexing (get a row, column, or batch)Slice indexing (get a subset of rows, columns, or batches)LongTensor indexing (get a set of individual entries by index)Ellipses (support indexing operations with arbitrary batch dimensions)D=DiagLinearOperator(torch.randn(2,3,100))# Represents an operator of size 2 x 3 x 100 x 100D[-1]# Returns a 3 x 100 x 100 operatorD[...,:10,-5:]# Returns a 2 x 3 x 10 x 5 operatorD[...,torch.LongTensor([0,1,2,3]),torch.LongTensor([0,1,2,3])]# Returns a 2 x 3 x 4 tensorComposition and DecorationLinearOperators can be composed with one another in various ways.
This includesAddition (LinearOpA + LinearOpB)Matrix multiplication (LinearOpA @ LinearOpB)Concatenation (torch.cat([LinearOpA, LinearOpB], dim=-2))Kronecker product (torch.kron(LinearOpA, LinearOpB))In addition, there are many ways to "decorate" LinearOperator objects.
This includes:Elementwise multiplying by constants (torch.mul(2., LinearOpA))Summing over batches (torch.sum(LinearOpA, dim=-3))Elementwise multiplying over batches (torch.prod(LinearOpA, dim=-3))See the documentation for afull list of supported composition and decoration operations.InstallationLinearOperator requires Python >= 3.8.Standard Installation (Most Recent Stable Version)We recommend installing viapipor Anaconda:pipinstalllinear_operator# orcondainstalllinear_operator-cgpytorchThe installation requires the following packages:PyTorch >= 1.11ScipyYou can customize your PyTorch installation (i.e. CUDA version, CPU only option)
by following thePyTorch installation instructions.Installing from themainBranch (Latest Unsable Version)To install what is currently on themainbranch (potentially buggy and unstable):pipinstall--upgradegit+https://github.com/cornellius-gp/linear_operator.gitDevelopment InstallationIf you are contributing a pull request, it is best to perform a manual installation:gitclonehttps://github.com/cornellius-gp/linear_operator.gitcdlinear_operator
pipinstall-e".[dev,docs,test]"ContributingSee the contributing guidelinesCONTRIBUTING.mdfor information on submitting issues and pull requests.LicenseLinearOperator isMIT licensed.