Package: torch-linops
An Abstract Linear Operator Library for pyTorchThis library implements a generic structure for abstract linear operators and
enables a number of standard operations on them:Arithmetic:A + B,A - B,-A,A @ Ball work exactly as expected to
combine linear operators.Indexing:A[k:ell,m:n]works as expected.Solves:Ax = bcan be solved withCGfor PSD matrices,minresfor
symmetric matrices,LSQR(to be implemented), orLSMR(to be implemented).Trace estimation: The trace of square matrices, can be estimated via Hutch++
and Hutchinson's estimator.Diamond-Boyd stochastic equilibrationRandomized Nyström PreconditioningAutomatic adjoint operator generation.UsingLinearOperatorsThe public API of theLinearOperatorlibrary is that everyLinearOperatorhas the
following properties and methods:classLinearOperator:# Propertiesshape:tuple[int,int]T:LinearOperatorsupports_operator_matrix:booldevice:torch.Device# Matrix multiplydef__matmul__(self,b:torch.Tensor)->torch.Tensor:...def__rmatmul__(self,b:torch.Tensor)->torch.Tensor:...def__matmul__(self,b:LinearOperator)->LinearOperator:...def__rmatmul__(self,b:LinearOperator)->LinearOperator:...# Linear Solve Methodsdefsolve_I_p_lambda_AT_A_x_eq_b(self,lambda_:float,b:torch.Tensor,x0:torch.Tensor|None=None,*,precondition:None|Literal['nsytrom'],hot=False)->torch.Tensor:...defsolve_A_x_eq_b(self,b:torch.Tensor,x0:torch.Tensor|None=None)->torch.Tensor:...# Transformations on LinearOperatordef__mul__(self,c:float)->LinearOperator:...def__rmul__(self,c:float)->LinearOperator:...def__truediv__(self,c:float)->LinearOperator:...def__pow__(self,k:int)->LinearOperator:...def__add__(self,c:LinearOperator)->LinearOperator:...def__sub__(self,c:LinearOperator)->LinearOperator:...def__neg__(self)->LinearOperator:...def__pos__(self)->LinearOperator:...def__getitem__(self,key)->LinearOperator:...The following functions are available in the root of the library:defoperator_matrix_product(A:LinearOperator,M:torch.Tensor)->torch.Tensor:...defaslinearoperator(A:torch.Tensor|LinearOperator)->LinearOperator:...defvstack(ops:list[LinearOperator]|tuple[LinearOperator,...])->LinearOperator:...defhstack(ops:list[LinearOperator]|tuple[LinearOperator,...])->LinearOperator:...# To be implemented:defbmat(ops:list[list[LinearOperator]])->LinearOperator:...# Optimizes out ZeroOperatorThe following functions are available inlinops.tracefor trace estimation:defhutchpp(A:lo.LinearOperator,m:int)->float:...defhutchinson(A:lo.LinearOperator,m:int)->float:...linops.equilibrationcontainsequilibrateandsymmetric_equilibrate.
Their public API is not finalized, if you wish to use them it is recommend you read the source code.Creating Linear OperatorsLinear operators can be constructed in the following way:Creating a sub-class ofLinearOperatorCalling one of the following constructors:IdentityOperator(n: int)DiagonalOperator(diag: torch.Tensor): wherediagis a 1D torch tensor.MatrixOperator(M: torch.Tensor): whereMis a 2D torch tensor.SelectionOperator(shape: tuple[int, int], idxs: slice | list[int | slice])KKTOperator(H: LinearOperator, A: LinearOperator): whereHis a squareLinearOperatorandAis aLinearOperatorVectorJacobianOperator(f: torch.Tensor, x: torch.Tensor): wherefis the output of the function being differentiated
which has a torch autograd value andxis the vector on whichensures_gradwas called.ZeroOperator(shape: tuple[int, int])Combining operators via:A + B,A - B,A @ BforA,Blinear operatorshstack,vstackA,c A,A / c,v * A,A / vfor scalarcand vectorv.ImplementingLinearOperatorsTo implement aLinearOperatorthe following are mandatory:Set_shape: tuple[int, int]to the shape of the operator.Setdeviceappropriately, if the operator requires vectors to be on a particular device.Implement a methoddef _matmul_impl(self, v: torch.Tensor) -> torch.Tensor: ...that implements your matrix vector product.The following are recommended to improve performance:If your_matmul_implmethod handles matrix inputs correctly, setsupports_operator_matrix: booltoTrue.If it is possible to describe the adjoint operator, set_adjoint: LinearOperatorto point to the adjoint of your operator. If you do not compute this, then one will be autogenerated by differentiating through your_matmul_impl.It is suggested that, if possible, you replace any other methods with specialized implementations.