fca285821740bcc013bfe27b5dd11b7fdb9b9812,gpytorch/lazy/kronecker_product_added_diag_lazy_tensor.py,KroneckerProductAddedDiagLazyTensor,_solve,#KroneckerProductAddedDiagLazyTensor#Any#Any#Any#,49

Before Change

                lt = self.lazy_tensor.to(torch.double)
                dlt = self.diag_tensor.to(torch.double)

                KDinv = KroneckerProductLazyTensor(
                    *[tfull.matmul(tdiag.inverse()) for tfull, tdiag in zip(lt.lazy_tensors, dlt.lazy_tensors)]
                )
                // TODO: Figure out how to cache the decompositon for use in later solves
                Lambda, S = KDinv.symeig(eigenvectors=True)
                LambdaI = DiagLazyTensor(Lambda + 1)
                tmp_term = S.matmul(LambdaI.inv_matmul(S._transpose_nonbatch().matmul(rhs)))
                res = lt._solve(rhs - tmp_term, preconditioner=preconditioner, num_tridiag=num_tridiag)
                return res.to(rhs_dtype)

After Change

                // Each sub-matrix D_i^{-1} has constant diagonal, so we may scale the eigenvalues of the
                // eigendecomposition of K_i by its inverse to get an eigendecomposition of K_i D_i^{-1}.
                sub_evals, sub_evecs = [], []
                for lt_, dlt_ in zip(lt.lazy_tensors, dlt.lazy_tensors):
                    evals_, evecs_ = lt_.symeig(eigenvectors=True)
                    sub_evals.append(DiagLazyTensor(evals_ / dlt_.diag_values))
                    sub_evecs.append(evecs_)
                Lambda_I = KroneckerProductDiagLazyTensor(*sub_evals).add_jitter(1.0)
                S = KroneckerProductLazyTensor(*sub_evecs)
                tmp_term = S.matmul(Lambda_I.inv_matmul(S._transpose_nonbatch().matmul(rhs)))
                res = lt._solve(rhs - tmp_term, preconditioner=preconditioner, num_tridiag=num_tridiag)
                return res.to(rhs_dtype)

Instances