Repository: cornellius-gp/gpytorch
Issue: [Bug] - New release broke SGPR? 

# 🐛 Bug

Hi and thanks a lot for your work. Since the new release when i am trying to predict with the SGPR model, i get an error message that says that it is likely a bug in GPyTorch.

Thanks in advance.

## To reproduce

```python
import numpy as np
from tslearn.datasets import UCR_UEA_datasets
import math
import torch
import gpytorch
from gpytorch.means import ConstantMean
from gpytorch.kernels import ScaleKernel, RBFKernel, MaternKernel, InducingPointKernel
from gpytorch.distributions import MultivariateNormal

# Make plots inline
%matplotlib inline

inducing_p = 0

from gpytorch.means import ConstantMean
from gpytorch.kernels import ScaleKernel, RBFKernel, InducingPointKernel
from gpytorch.distributions import MultivariateNormal

class GPRegressionModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
        global inducing_p
        self.mean_module = ConstantMean()
        self.base_covar_module = RBFKernel()
        self.covar_module = InducingPointKernel(self.base_covar_module, inducing_points=torch.linspace(torch.min(train_x), torch.max(train_x), inducing_p), likelihood=likelihood)

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return MultivariateNormal(mean_x, covar_x)

#global inducing_p
X_train, y_train, X_test, y_test = UCR_UEA_datasets().load_dataset('Beef')

data = np.r_[X_train, X_test]
train_y = X_train[0,:,0]  

for indx in range(1, 20):
    train_y = np.vstack([train_y, data[indx,:,0]])

train_x = torch.from_numpy(np.arange(1,train_y.shape[1]+1)).float()
train_y = torch.from_numpy(train_y).float()
train_x = (train_x - train_x.mean(0)) / train_x.std(0)

#Normalization [-1,1]
for i in range(20):
    train_y[i] = train_y[i] - torch.min(train_y[i])
    train_y[i] = 2 * (train_y[i] / torch.max(train_y[i])) - 1  

inducing_p = 2*(int(math.log2(train_y.shape[1]))+1)
# initialize likelihood and model
likelihood = []
model = []
for i in range(20):
    likelihood.append(gpytorch.likelihoods.GaussianLikelihood())
    model.append(GPRegressionModel(train_x, train_y[i], likelihood[i]))

training_iter = 50
for i in range(20):
    # Find optimal model hyperparameters
    model[i].train()
    likelihood[i].train()

    # Use the adam optimizer
    optimizer = torch.optim.Adam(model[i].parameters(), lr=0.1)  # Includes GaussianLikelihood parameters
    # "Loss" for GPs - the marginal log likelihood
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood[i], model[i])

    for it in range(training_iter):
        # Zero gradients from previous iteration
        optimizer.zero_grad()
        # Output from model
        output = model[i](train_x)
        # Calc loss and backprop gradients
        loss = -mll(output, train_y[i])
        loss.backward()        
        optimizer.step()


# Get into evaluation (predictive posterior) mode
for i in range(20):
    model[i].eval()
    likelihood[i].eval()

# Test points are regularly spaced along [0,1]
# Make predictions by feeding model through likelihood
observed_pred = []
cross_covar = []
preds = []
with gpytorch.settings.max_preconditioner_size(10), torch.no_grad():
    for i in range(20):
        preds.append(model[i](train_x))               
```
## Error message
```
ValueError                                Traceback (most recent call last)
<ipython-input-20-e9aa2e04e88d> in <module>
     90 with gpytorch.settings.max_preconditioner_size(10), torch.no_grad():
     91     for i in range(20):
---> 92         preds.append(model[i](train_x))

~/miniconda3/lib/python3.8/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
    317             # Make the prediction
    318             with settings._use_eval_tolerance():
--> 319                 predictive_mean, predictive_covar = self.prediction_strategy.exact_prediction(full_mean, full_covar)
    320 
    321             # Reshape predictive mean to match the appropriate event shape

~/miniconda3/lib/python3.8/site-packages/gpytorch/models/exact_prediction_strategies.py in exact_prediction(self, joint_mean, joint_covar)
    805         return (
    806             self.exact_predictive_mean(test_mean, test_train_covar),
--> 807             self.exact_predictive_covar(test_test_covar, test_train_covar),
    808         )
    809 

~/miniconda3/lib/python3.8/site-packages/gpytorch/models/exact_prediction_strategies.py in exact_predictive_covar(self, test_test_covar, test_train_covar)
    815         if not isinstance(test_train_covar, MatmulLazyTensor):
    816             # We should not hit this point of the code - this is to catch potential bugs in GPyTorch
--> 817             raise ValueError(
    818                 f"Expected SGPR output to be a MatmulLazyTensor. Got {test_train_covar.__class__.__name__} instead. "
    819                 "This is likely a bug in GPyTorch."

ValueError: Expected SGPR output to be a MatmulLazyTensor. Got LowRankRootAddedDiagLazyTensor instead. This is likely a bug in GPyTorch.
```

## System information

- GPyTorch: Version 1.4.1
- PyTorch Version: 1.8.1+cu102
- Computer OS: Ubuntu 18.04.5 LTS

Changed files:
test/examples/test_sgpr_regression.py, gpytorch/models/exact_prediction_strategies.py

diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
index 93e11a4e..b220ddf5 100644
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -12,6 +12,7 @@ from ..lazy import (
     ConstantMulLazyTensor,
     InterpolatedLazyTensor,
     LazyEvaluatedKernelTensor,
+    LowRankRootAddedDiagLazyTensor,
     MatmulLazyTensor,
     NonLazyTensor,
     RootLazyTensor,
@@ -812,13 +813,19 @@ class SGPRPredictionStrategy(DefaultPredictionStrategy):
         # covar_cache = K_{UU}^{-1/2} K_{UX}( K_{XX} + \\sigma^2 I )^{-1} K_{XU} K_{UU}^{-1/2}
 
         # Decompose test_train_covar = l, r
-        if not isinstance(test_train_covar, MatmulLazyTensor):
+        # Main case: test_x and train_x are different - test_train_covar is a MatmulLazyTensor
+        if isinstance(test_train_covar, MatmulLazyTensor):
+            L = test_train_covar.left_lazy_tensor.evaluate()
+        # Edge case: test_x and train_x are the same - test_train_covar is a LowRankRootAddedDiagLazyTensor
+        elif isinstance(test_train_covar, LowRankRootAddedDiagLazyTensor):
+            L = test_train_covar._lazy_tensor.root.evaluate()
+        else:
             # We should not hit this point of the code - this is to catch potential bugs in GPyTorch
             raise ValueError(
-                f"Expected SGPR output to be a MatmulLazyTensor. Got {test_train_covar.__class__.__name__} instead. "
+                "Expected SGPR output to be a MatmulLazyTensor or AddedDiagLazyTensor. "
+                f"Got {test_train_covar.__class__.__name__} instead. "
                 "This is likely a bug in GPyTorch."
             )
-        L = test_train_covar.left_lazy_tensor.evaluate()
 
         res = test_test_covar - (L @ (covar_cache @ L.transpose(-1, -2)))
         return res
diff --git a/test/examples/test_sgpr_regression.py b/test/examples/test_sgpr_regression.py
index 2fb450ed..3420aa18 100644
--- a/test/examples/test_sgpr_regression.py
+++ b/test/examples/test_sgpr_regression.py
@@ -101,17 +101,22 @@ class TestSGPRRegression(unittest.TestCase, BaseTestCase):
         test_preds = likelihood(gp_model(test_x)).mean
         mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
 
-        self.assertLess(mean_abs_error.squeeze().item(), 0.05)
+        self.assertLess(mean_abs_error.squeeze().item(), 0.1)
 
         # Test variances
         test_vars = likelihood(gp_model(test_x)).variance
         self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
-        self.assertGreater(test_vars.min().item() + 0.05, likelihood.noise.item())
+        self.assertGreater(test_vars.min().item() + 0.1, likelihood.noise.item())
         self.assertLess(
             test_vars.max().item() - 0.05,
             likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
         )
 
+        # Test on training data
+        test_outputs = likelihood(gp_model(train_x))
+        self.assertLess((test_outputs.mean - train_y).max().item(), 0.1)
+        self.assertLess(test_outputs.variance.max().item(), likelihood.noise.item() * 2)
+
     def test_sgpr_mean_abs_error_cuda(self):
         # Suppress numerical warnings
         warnings.simplefilter("ignore", NumericalWarning)
@@ -134,7 +139,7 @@ class TestSGPRRegression(unittest.TestCase, BaseTestCase):
             # Test variances before optimization
             test_vars = likelihood(gp_model(test_x)).variance
             self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
-            self.assertGreater(test_vars.min().item() + 0.05, likelihood.noise.item())
+            self.assertGreater(test_vars.min().item() + 0.1, likelihood.noise.item())
             self.assertLess(
                 test_vars.max().item() - 0.05,
                 likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
@@ -169,7 +174,7 @@ class TestSGPRRegression(unittest.TestCase, BaseTestCase):
             # Test variances
             test_vars = likelihood(gp_model(test_x)).variance
             self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
-            self.assertGreater(test_vars.min().item() + 0.05, likelihood.noise.item())
+            self.assertGreater(test_vars.min().item() + 0.1, likelihood.noise.item())
             self.assertLess(
                 test_vars.max().item() - 0.05,
                 likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()