Repository: cornellius-gp/gpytorch

Issue:
Odd prediction behavior when Mixing Sum and Product Kernels
Hi,

So I think this is worth some attention. I'm using a kernel that is a product of two kernels, one of which is a sum kernel. So, some fake data first...

```
import numpy as np
import pandas as pd
import gpytorch
import torch
torch.set_default_tensor_type('torch.DoubleTensor')

import matplotlib.pyplot as plt
%matplotlib inline

x = np.linspace(0,1,500)
beta = 10
offset = 2
sin_amplitutde = 3
sin_period = .1
noise_var = 1

y = x*beta + offset + np.sin(x*2*np.pi/sin_period)*sin_amplitutde + np.random.normal(0,np.sqrt(noise_var),len(x))
y = 1.5*x*y

def clean_ax(ax):
    ax.grid(True)
    for which in ['top','right']:
        ax.spines[which].set_visible(False)

figsize=(10,6)
fig, ax = plt.subplots(figsize=figsize)
ax.plot(x,y,'*')
clean_ax(ax)
```
![image](https://user-images.githubusercontent.com/19956442/53211601-9fc18e80-35f6-11e9-8b28-06c339fe435e.png)

What I suspect should work is:

```
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        
        ker1 = gpytorch.kernels.LinearKernel(num_dimensions=1, active_dims=[0])
        
        ker2 = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(active_dims=[1]))
        ker3 = gpytorch.kernels.LinearKernel(num_dimensions=1, active_dims=[0])
        
        prodker = gpytorch.kernels.LinearKernel(num_dimensions=1, active_dims=[0])
        
        # This line below is what I think is worthy of special attention.
        self.covar_module = prodker*(ker1 + ker2 + ker3) 
#         self.covar_module = (ker1 + ker2 + ker3) #If we do this, we don't have an issue. 
    
    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 = np.concatenate([x.reshape(-1,1),np.mod(x,sin_period).reshape(-1,1)],axis=1)
X = torch.from_numpy(X).to(torch.float64)
y = torch.from_numpy(y).to(torch.float64)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = ExactGPModel(X, y, likelihood)
```

Now train...

```
model.train()
likelihood.train()

optimizer = torch.optim.Adam([
    {'params': model.parameters()},
], lr=0.1)

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

training_iter = 200
for i in range(training_iter):
    # Zero gradients from previous iteration
    optimizer.zero_grad()
    # Output from model
    output = model(X)
    # Calc loss and backprop gradients
    loss = -mll(output, y)
    loss.backward()
    if not i % 20:
        print('i: {0}, Loss: {1}'.format(i, loss.item()))
    optimizer.step()
```

Great, and this works:

```
model.eval()
likelihood.eval()

with torch.no_grad(), gpytorch.settings.fast_pred_var():
    y_pred = likelihood(model(X)).mean.numpy()

figsize=(10,6)
fig, ax = plt.subplots(figsize=figsize)
ax.plot(x,y.numpy(),'*')
clean_ax(ax)
ax.plot(x,y_pred)
```

![image](https://user-images.githubusercontent.com/19956442/53211651-d7303b00-35f6-11e9-8b50-77fb75d7ee45.png)

Nice! But if I slice `X`, it breaks:

```
model.eval()
likelihood.eval()

with torch.no_grad(), gpytorch.settings.fast_pred_var():
    y_pred = likelihood(model(X[:50,:])).mean.numpy()
```
 With this error:

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-67-d8398cf67836> in <module>
      3 
      4 with torch.no_grad(), gpytorch.settings.fast_pred_var():
----> 5     y_pred = likelihood(model(X[:50,:])).mean.numpy()

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
    261             test_train_covar = full_covar[..., num_train:, :num_train]
    262 
--> 263             predictive_mean = self.prediction_strategy.exact_predictive_mean(test_mean, test_train_covar)
    264             predictive_covar = self.prediction_strategy.exact_predictive_covar(test_test_covar, test_train_covar)
    265 

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/models/exact_prediction_strategies.py in exact_predictive_mean(self, test_mean, test_train_covar)
    289             if self.non_batch_train and precomputed_cache.dim() == 2:
    290                 precomputed_cache = precomputed_cache[0]
--> 291             res = test_train_covar.matmul(precomputed_cache)
    292 
    293         res = res + test_mean

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in matmul(self, other)
    804             return MatmulLazyTensor(self, other)
    805 
--> 806         func = Matmul(self.representation_tree())
    807         return func(other, *self.representation())
    808 

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in representation_tree(self)
    236 
    237     def representation_tree(self):
--> 238         return LazyTensorRepresentationTree(self.evaluate_kernel())

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
    197             with settings.lazily_evaluate_kernels(False):
    198                 self._cached_kernel_eval = self.kernel(
--> 199                     x1, x2, diag=False, batch_dims=self.batch_dims, **self.params
    200                 )
    201             if self.squeeze_row:

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/kernels/kernel.py in __call__(self, x1, x2, diag, batch_dims, **params)
    382                 res = LazyEvaluatedKernelTensor(self, x1_, x2_, batch_dims=batch_dims, **params)
    383             else:
--> 384                 res = super(Kernel, self).__call__(x1_, x2_, batch_dims=batch_dims, **params)
    385 
    386             # Now we'll make sure that the shape we're getting makes sense

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     18 
     19     def __call__(self, *inputs, **kwargs):
---> 20         outputs = self.forward(*inputs, **kwargs)
     21         if isinstance(outputs, list):
     22             return [_validate_module_outputs(output) for output in outputs]

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/kernels/kernel.py in forward(self, x1, x2, **params)
    460         for kern in self.kernels[1:]:
    461             next_term = kern(x1, x2, **params)
--> 462             res = res * lazify(next_term)
    463         return res
    464 

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in __mul__(self, other)
   1372             return other * self
   1373 
-> 1374         return self.mul(other)
   1375 
   1376     def __setattr__(self, name, val):

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in mul(self, other)
    838             from .mul_lazy_tensor import MulLazyTensor
    839 
--> 840             return MulLazyTensor(self, other).evaluate_kernel()
    841 
    842         else:

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in evaluate_kernel(self)
    633         all lazily evaluated kernels actually evaluated.
    634         """
--> 635         return self.representation_tree()(*self.representation())
    636 
    637     def inv_matmul(self, right_tensor, left_tensor=None):

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/mul_lazy_tensor.py in representation_tree(self)
    238 
    239     def representation_tree(self):
--> 240         if self.non_lazy_self is not None:
    241             return self.non_lazy_self.representation_tree()
    242         else:

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/mul_lazy_tensor.py in non_lazy_self(self)
     37         if hasattr(self, "_non_lazy_self"):
     38             return self._non_lazy_self[0]
---> 39         elif len(self._args) == 1:
     40             return self._args[0]
     41         else:

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/mul_lazy_tensor.py in _args(self)
     89                 # Choose which we're doing: root decomposition or exact
     90                 if left_lazy_tensor.root_decomposition_size() < left_lazy_tensor.size(-1):
---> 91                     left_lazy_tensor = left_lazy_tensor.root_decomposition()
     92                     right_lazy_tensor = right_lazy_tensor.root_decomposition()
     93 

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
     15         cache_name = name if name is not None else method
     16         if cache_name not in self._memoize_cache:
---> 17             self._memoize_cache[cache_name] = method(self, *args, **kwargs)
     18         return self._memoize_cache[cache_name]
     19 

~/.virtualenvs/GenRS/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in root_decomposition(self)
   1019             raise RuntimeError(
   1020                 "root_decomposition only operates on (batches of) square (symmetric) LazyTensors. "
-> 1021                 "Got a {} of size {}.".format(self.__class__.__name__, self.size())
   1022             )
   1023 

RuntimeError: root_decomposition only operates on (batches of) square (symmetric) LazyTensors. Got a SumLazyTensor of size torch.Size([1, 50, 500]).
```

`X` having a different number of rows than what was fed as training data seems to be the issue. This looks like a bug to me. Thought you'd want to know. It seems to be specific to the product/sum thing as well. If i change `self.covar_module = prodker*(ker1 + ker2 + ker3)`  --> `self.covar_module = ker1 + ker2 + ker3`, we don't have this issue.


Discussion:
This occurred because LinearKernel actually uses a non-trivial LazyTensor in the back end for efficient Bayesian linear regression, which caused the whole composition to create a MulLazyTensor even in the rectangular `K_{x*,X}` case.

#528 resolves this, and will also make certain kernel compositions much more efficient to work with.

Resolving patch:
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -3,7 +3,7 @@
 from abc import abstractmethod
 import torch
 from torch.nn import ModuleList
-from ..lazy import lazify, LazyEvaluatedKernelTensor, ZeroLazyTensor
+from ..lazy import lazify, delazify, LazyEvaluatedKernelTensor, ZeroLazyTensor
 from ..module import Module
 from .. import settings
 from ..utils.transforms import _get_inv_param_transform
@@ -456,10 +456,21 @@ def __init__(self, *kernels):
         self.kernels = ModuleList(kernels)
 
     def forward(self, x1, x2, **params):
-        res = lazify(self.kernels[0](x1, x2, **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))
+        else:
+            res = lazify(self.kernels[0](x1, x2, **params))
+
         for kern in self.kernels[1:]:
             next_term = kern(x1, x2, **params)
-            res = res * lazify(next_term)
+            if not x1_eq_x2:
+                # Again delazify if x1 != x2
+                res = res * delazify(next_term)
+            else:
+                res = res * lazify(next_term)
         return res
 
     def size(self, x1, x2):
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -827,6 +827,14 @@ def mul(self, other):
             :obj:`gpytorch.lazy.ConstantMulLazyTensor`. If other was
             another matrix, this will likely be a :obj:`gpytorch.lazy.MulLazyTensor`.
         """
+        from .zero_lazy_tensor import ZeroLazyTensor
+        from .diag_lazy_tensor import DiagLazyTensor
+
+        if isinstance(other, ZeroLazyTensor):
+            return other
+        elif isinstance(other, DiagLazyTensor):
+            return other * self
+
         if not (torch.is_tensor(other) or isinstance(other, LazyTensor)) or (
             torch.is_tensor(other) and (other.numel() == 1 or (self.dim() == 3 and other.numel() == self.size(0)))
         ):
@@ -1363,14 +1371,6 @@ def __mul__(self, other):
         Convenience alias of :meth:`~gpytorch.lazy.LazyTensor.mul` that allows the standard product operator to be
         used.
         """
-        from .zero_lazy_tensor import ZeroLazyTensor
-        from .diag_lazy_tensor import DiagLazyTensor
-
-        if isinstance(other, ZeroLazyTensor):
-            return other
-        elif isinstance(other, DiagLazyTensor):
-            return other * self
-
         return self.mul(other)
 
     def __setattr__(self, name, val):
diff --git a/gpytorch/lazy/non_lazy_tensor.py b/gpytorch/lazy/non_lazy_tensor.py
--- a/gpytorch/lazy/non_lazy_tensor.py
+++ b/gpytorch/lazy/non_lazy_tensor.py
@@ -62,6 +62,18 @@ def diag(self):
     def evaluate(self):
         return self.tensor
 
+    def __add__(self, other):
+        if isinstance(other, NonLazyTensor):
+            return NonLazyTensor(self.tensor + other.tensor)
+        else:
+            return super(NonLazyTensor, self).__add__(other)
+
+    def mul(self, other):
+        if isinstance(other, NonLazyTensor):
+            return NonLazyTensor(self.tensor * other.tensor)
+        else:
+            return super(NonLazyTensor, self).mul(other)
+
 
 def lazify(obj):
     """


Test patch: