InducingPointKernel uses LowRankRootLazyTensor
@@ -7,7 +7,7 @@ import torch
from .. import settings
from ..distributions import MultivariateNormal
-from ..lazy import DiagLazyTensor, MatmulLazyTensor, PsdSumLazyTensor, RootLazyTensor, delazify
+from ..lazy import DiagLazyTensor, LowRankRootAddedDiagLazyTensor, LowRankRootLazyTensor, MatmulLazyTensor, delazify
from ..mlls import InducingPointKernelAddedLossTerm
from ..utils.cholesky import psd_safe_cholesky
from .kernel import Kernel
@@ -56,11 +56,11 @@ class InducingPointKernel(Kernel):
def _get_covariance(self, x1, x2):
k_ux1 = delazify(self.base_kernel(x1, self.inducing_points))
if torch.equal(x1, x2):
- covar = RootLazyTensor(k_ux1.matmul(self._inducing_inv_root))
+ covar = LowRankRootLazyTensor(k_ux1.matmul(self._inducing_inv_root))
# Diagonal correction for predictive posterior
correction = (self.base_kernel(x1, x2, diag=True) - covar.diag()).clamp(0, math.inf)
- covar = PsdSumLazyTensor(covar, DiagLazyTensor(correction))
+ covar = LowRankRootAddedDiagLazyTensor(covar, DiagLazyTensor(correction))
else:
k_ux2 = delazify(self.base_kernel(x2, self.inducing_points))
covar = MatmulLazyTensor(
