Repository: cornellius-gp/gpytorch (python)

Issue:
ValueError: NonLazyTensor expects a matrix (or batches of matrices) - got a Tensor of size torch.Size([50]).
# 🐛 Bug

Hi,
I encounter an error when running SVGP with an additive kernel, which did not happen to me when using gpytorch 0.2.1.
```
~/anaconda3/envs/gpytorch2/lib/python3.6/site-packages/gpytorch/variational/whitened_variational_strategy.py in forward(self, x)
    188 
    189             if self.training:
--> 190                 data_covariance = DiagLazyTensor((data_data_covar.diag() - interp_data_data_var).clamp(0, math.inf))
    191             else:
    192                 neg_induc_data_data_covar = induc_induc_covar.inv_matmul(
```

** Code snippet to reproduce **
```python
class SVGPRegressionModel(gpytorch.models.AbstractVariationalGP):
    def __init__(self, inducing_points):
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(10)
        variational_strategy = gpytorch.variational.WhitenedVariationalStrategy(self, inducing_points, variational_distribution, learn_inducing_locations=True)
        super(SVGPRegressionModel, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel() + gpytorch.kernels.MaternKernel(nu=2.5))

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

x = torch.randn(50, 2)
z = torch.randn(10, 2)
model = SVGPRegressionModel(z)
output = model(x)
```


Pull request description:
Fix composition kernels in diag mode
On master, the composition kernels (`AdditiveKernel`, `ProductKernel`) don't handle `diag` mode properly anymore with the kernel `__call__` changes in `0.3.1`. This PR fixes that, and adds unit tests to ensure that we catch this in the future.

Fixes #655 

Patch:
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 0b256c38..69ca7887 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -472,11 +472,15 @@ class AdditiveKernel(Kernel):
         super(AdditiveKernel, self).__init__()
         self.kernels = ModuleList(kernels)
 
-    def forward(self, x1, x2, **params):
-        res = ZeroLazyTensor()
+    def forward(self, x1, x2, diag=False, **params):
+        res = ZeroLazyTensor() if not diag else 0
         for kern in self.kernels:
-            next_term = kern(x1, x2, **params)
-            res = res + lazify(next_term)
+            next_term = kern(x1, x2, diag=diag, **params)
+            if not diag:
+                res = res + lazify(next_term)
+            else:
+                res = res + next_term
+
         return res
 
     def prediction_strategy(self, train_inputs, train_prior_dist, train_labels, likelihood):
@@ -500,22 +504,29 @@ class ProductKernel(Kernel):
         super(ProductKernel, self).__init__()
         self.kernels = ModuleList(kernels)
 
-    def forward(self, x1, x2, **params):
+    def forward(self, x1, x2, diag=False, **params):
         x1_eq_x2 = torch.equal(x1, x2)
 
         if not x1_eq_x2:
             # If x1 != x2, then we can't make a MulLazyTensor because the kernel won't necessarily be square/symmetric
-            res = delazify(self.kernels[0](x1, x2, **params))
+            res = delazify(self.kernels[0](x1, x2, diag=diag, **params))
         else:
-            res = lazify(self.kernels[0](x1, x2, **params))
+            res = self.kernels[0](x1, x2, diag=diag, **params)
+
+            if not diag:
+                res = lazify(res)
 
         for kern in self.kernels[1:]:
-            next_term = kern(x1, x2, **params)
+            next_term = kern(x1, x2, diag=diag, **params)
             if not x1_eq_x2:
                 # Again delazify if x1 != x2
                 res = res * delazify(next_term)
             else:
-                res = res * lazify(next_term)
+                if not diag:
+                    res = res * lazify(next_term)
+                else:
+                    res = res * next_term
+
         return res
 
     def num_outputs_per_input(self, x1, x2):


Test patch:
diff --git a/test/kernels/test_additive_kernel.py b/test/kernels/test_additive_kernel.py
index 301ac94c..5c6b1515 100644
--- a/test/kernels/test_additive_kernel.py
+++ b/test/kernels/test_additive_kernel.py
@@ -39,6 +39,75 @@ class TestAdditiveKernel(unittest.TestCase):
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 2e-5)
 
+    def test_computes_sum_of_radial_basis_function_diag(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2, 2], dtype=torch.float).view(3, 1)
+
+        actual = torch.tensor([0.2702, 2.000, 0.0222])
+
+        lengthscale = 2.
+
+        kernel_1 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel_2 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel = kernel_1 + kernel_2
+        kernel.eval()
+
+        kernel.eval()
+        res = kernel(a, b, diag=True)
+        self.assertLess(torch.norm(res - actual), 1e-3)
+
+    def test_computes_sum_of_three_radial_basis_function_diag(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2, 2], dtype=torch.float).view(3, 1)
+
+        actual = torch.tensor([0.4060, 3.000, 0.0333])
+
+        lengthscale = 2.
+
+        kernel_1 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel_2 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel_3 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel = kernel_1 + kernel_2 + kernel_3
+        kernel.eval()
+
+        res = kernel(a, b, diag=True)
+        self.assertLess(torch.norm(res - actual), 1e-3)
+
+    def test_computes_product_of_radial_basis_function_diag(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2, 2], dtype=torch.float).view(3, 1)
+
+        actual = torch.tensor([2.4788e-03, 1.000, 1.3710e-06])
+
+        lengthscale = 2.
+
+        kernel_1 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel_2 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel_3 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel = kernel_1 * kernel_2 * kernel_3
+        kernel.eval()
+
+        kernel.eval()
+        res = kernel(a, b, diag=True)
+        self.assertLess(torch.norm(res - actual), 1e-3)
+
+    def test_computes_product_of_three_radial_basis_function_diag(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2, 2], dtype=torch.float).view(3, 1)
+
+        actual = torch.tensor([1.8316e-02, 1.000, 1.2341e-04])
+
+        lengthscale = 2.
+
+        kernel_1 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel_2 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel = kernel_1 * kernel_2
+        kernel.eval()
+
+        kernel.eval()
+        res = kernel(a, b, diag=True)
+        self.assertLess(torch.norm(res - actual), 1e-3)
+
     def test_computes_product_of_three_radial_basis_function(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
         b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)