Repository: cornellius-gp/gpytorch

Issue:
[Bug] SM + GridInterpolationKernel Interoperability
# 🐛 Bug

Looks like the updates to spectral mixture kernels have broken its interoperability with GridInterpolation Kernels.


## To reproduce

** Code snippet to reproduce **
```python
from gpytorch.kernels import SpectralMixtureKernel, GridInterpolationKernel

covar_module = SpectralMixtureKernel(num_mixtures=3)
grid_interp = GridInterpolationKernel(covar_module, grid_size=12, grid_bounds = [[0, 1]])

grid_interp._inducing_forward(last_dim_is_batch=False)
# should be KroneckerProductLT
grid_interp._inducing_forward(last_dim_is_batch=True)
# should be ToeplitzLT
```

** Stack trace/error message **
```
~/Documents/GitHub/wjm_gpytorch/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
     57         kwargs_pkl = pickle.dumps(kwargs)
     58         if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):
---> 59             return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
     60         return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)
     61 

~/Documents/GitHub/wjm_gpytorch/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
    279             if res.shape != self.shape:
    280                 raise RuntimeError(
--> 281                     f"The expected shape of the kernel was {self.shape}, but got {res.shape}. "
    282                     "This is likely a bug in GPyTorch."
    283                 )

RuntimeError: The expected shape of the kernel was torch.Size([1, 1, 12]), but got torch.Size([3, 1, 12]). This is likely a bug in GPyTorch.
```

## Expected Behavior

Shouldn't have bugged out.

## System information

**Please complete the following information:**
gpytorch master
pytorch 1.7.0
Mac OS

## Additional context
Rolling back to aab81213 gives the expected output, but I haven't had a chance to debug more fully yet. 



Discussion:
Hmm okay I'll take a look

Resolving patch:
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -326,30 +326,26 @@ def forward(
         x2_cos = x2_ * self.mixture_means
 
         # Create grids
-        x1_exp_, x2_exp_ = self._create_input_grid(
-            x1_exp, x2_exp, diag=diag, last_dim_is_batch=last_dim_is_batch, **params
-        )
-        x1_cos_, x2_cos_ = self._create_input_grid(
-            x1_cos, x2_cos, diag=diag, last_dim_is_batch=last_dim_is_batch, **params
-        )
+        x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, diag=diag, **params)
+        x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, diag=diag, **params)
 
         # Compute the exponential and cosine terms
         exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi ** 2)
         cos_term = (x1_cos_ - x2_cos_).mul_(2 * math.pi)
         res = exp_term.exp_() * cos_term.cos_()
 
+        # Sum over mixtures
+        mixture_weights = self.mixture_weights.view(*self.mixture_weights.shape, 1, 1)
+        if not diag:
+            mixture_weights = mixture_weights.unsqueeze(-2)
+
+        res = (res * mixture_weights).sum(-3 if diag else -4)
+
         # Product over dimensions
         if last_dim_is_batch:
-            res = res.squeeze(-1)
+            # Put feature-dimension in front of data1/data2 dimensions
+            res = res.permute(*list(range(0, res.dim() - 3)), -1, -3, -2)
         else:
             res = res.prod(-1)
 
-        # Sum over mixtures
-        mixture_weights = self.mixture_weights.unsqueeze(-1)
-        if not diag:
-            mixture_weights = mixture_weights.unsqueeze(-1)
-        if last_dim_is_batch:
-            mixture_weights = mixture_weights.unsqueeze(-1)
-
-        res = (res * mixture_weights).sum(-2 if diag else -3)
         return res


Test patch: