Repository: cornellius-gp/gpytorch

Issue:
[Bug] IF-construction leads to RuntimeError in polynomial kernel 
# 🐛 Bug

I wanted to train a gp-model, so I made everything like in [tutorials](https://docs.gpytorch.ai/en/stable/examples/04_Variational_and_Approximate_GPs/SVGP_Multitask_GP_Regression.html). The code works with Linear and RBF kernels, but it does not work with Polynomial kernel.
It crushes with RuntimeError while inference. And from what I see in the code there might be several reasons why:
- misprint (if and else should be vice versa)
- pytorch changed the behavior of `torch.addmm`
- or maybe there is smth I do not understand :)

Please, look at the text below.

## To reproduce

** Code snippet to reproduce **
```python
import torch
import gpytorch

class MultitaskGPModel(gpytorch.models.ApproximateGP):
    def __init__(self, grid_size=4, num_latents=2, grid_bounds=[(0., 31.)], num_tasks=256):
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            grid_size, batch_shape=torch.Size([num_latents])
        )
        variational_strategy = gpytorch.variational.LMCVariationalStrategy(
            gpytorch.variational.GridInterpolationVariationalStrategy(self, grid_size, grid_bounds, variational_distribution),
            num_tasks=num_tasks,
            num_latents=num_latents,
            latent_dim=-1
        )
        super().__init__(variational_strategy)

        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_latents]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.PolynomialKernel(batch_shape=torch.Size([num_latents]), power=2),
            batch_shape=torch.Size([num_latents])
        )
        
    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
    
gp_layer = MultitaskGPModel()
likelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=256)

x = torch.randint(32, size=(1,), dtype=torch.float32).unsqueeze(1)

noise = gp_layer(x)
noise = likelihood(noise)
noise = noise.rsample((2, 2))
```

** Stack trace/error message **
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_15909/3264522926.py in <module>
     31 x = torch.randint(32, size=(1,), dtype=torch.float32).unsqueeze(1)
     32 
---> 33 noise = gp_layer(x)
     34 noise = likelihood(noise)
     35 noise = noise.rsample((2, 2))

~/anaconda3/lib/python3.8/site-packages/gpytorch/models/approximate_gp.py in __call__(self, inputs, prior, **kwargs)
     79         if inputs.dim() == 1:
     80             inputs = inputs.unsqueeze(-1)
---> 81         return self.variational_strategy(inputs, prior=prior, **kwargs)

~/anaconda3/lib/python3.8/site-packages/gpytorch/variational/lmc_variational_strategy.py in __call__(self, x, task_indices, prior, **kwargs)
    186             or ~gpytorch.distributions.MultivariateNormal (... x N)
    187         """
--> 188         latent_dist = self.base_variational_strategy(x, prior=prior, **kwargs)
    189         num_batch = len(latent_dist.batch_shape)
    190         latent_dim = num_batch + self.latent_dim

~/anaconda3/lib/python3.8/site-packages/gpytorch/variational/_variational_strategy.py in __call__(self, x, prior, **kwargs)
    108         # (Maybe) initialize variational distribution
    109         if not self.variational_params_initialized.item():
--> 110             prior_dist = self.prior_distribution
    111             self._variational_distribution.initialize_variational_distribution(prior_dist)
    112             self.variational_params_initialized.fill_(1)

~/anaconda3/lib/python3.8/site-packages/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 

~/anaconda3/lib/python3.8/site-packages/gpytorch/variational/grid_interpolation_variational_strategy.py in prior_distribution(self)
     74     def prior_distribution(self):
     75         out = self.model.forward(self.inducing_points)
---> 76         res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter())
     77         return res
     78 

~/anaconda3/lib/python3.8/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in add_jitter(self, jitter_val)
    275 
    276     def add_jitter(self, jitter_val=1e-3):
--> 277         return self.evaluate_kernel().add_jitter(jitter_val)
    278 
    279     def _unsqueeze_batch(self, dim):

~/anaconda3/lib/python3.8/site-packages/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 

~/anaconda3/lib/python3.8/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
    330             temp_active_dims = self.kernel.active_dims
    331             self.kernel.active_dims = None
--> 332             res = self.kernel(
    333                 x1,
    334                 x2,

~/anaconda3/lib/python3.8/site-packages/gpytorch/kernels/kernel.py in __call__(self, x1, x2, diag, last_dim_is_batch, **params)
    400                 res = LazyEvaluatedKernelTensor(x1_, x2_, kernel=self, last_dim_is_batch=last_dim_is_batch, **params)
    401             else:
--> 402                 res = lazify(super(Kernel, self).__call__(x1_, x2_, last_dim_is_batch=last_dim_is_batch, **params))
    403             return res
    404 

~/anaconda3/lib/python3.8/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     28 
     29     def __call__(self, *inputs, **kwargs):
---> 30         outputs = self.forward(*inputs, **kwargs)
     31         if isinstance(outputs, list):
     32             return [_validate_module_outputs(output) for output in outputs]

~/anaconda3/lib/python3.8/site-packages/gpytorch/kernels/scale_kernel.py in forward(self, x1, x2, last_dim_is_batch, diag, **params)
    101 
    102     def forward(self, x1, x2, last_dim_is_batch=False, diag=False, **params):
--> 103         orig_output = self.base_kernel.forward(x1, x2, diag=diag, last_dim_is_batch=last_dim_is_batch, **params)
    104         outputscales = self.outputscale
    105         if last_dim_is_batch:

~/anaconda3/lib/python3.8/site-packages/gpytorch/kernels/polynomial_kernel.py in forward(self, x1, x2, diag, last_dim_is_batch, **params)
     95 
     96         if x1.dim() == 2 and x2.dim() == 2:
---> 97             return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)
     98         else:
     99             return (torch.matmul(x1, x2.transpose(-2, -1)) + offset).pow(self.power)

RuntimeError: expand(torch.FloatTensor{[2, 1, 1]}, size=[4, 4]): the number of sizes provided (2) must be greater or equal to the number of dimensions in the tensor (3)
```

## System information

GPyTorch Version: 1.6.0
PyTorch Version : 1.11.0+cu113
Computer OS: Ubuntu 18.04.6 LTS 

## Additional context
I see in [`polynomial_kernel.py`](https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/kernels/polynomial_kernel.py) this piece of code:

```python
offset = self.offset.view(*self.batch_shape, 1, 1)
...
if x1.dim() == 2 and x2.dim() == 2:
    return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)
```

So the tensor `offset` has at least 3 dimensions or more, while `x1` and `x2` have exactly 2 dimensions, but `torch.addmm` will not work with such sizes. I do not know why such construction is needed, because the code from else-part would do the same but without an error.


Discussion:

Resolving patch:
diff --git a/gpytorch/kernels/polynomial_kernel.py b/gpytorch/kernels/polynomial_kernel.py
--- a/gpytorch/kernels/polynomial_kernel.py
+++ b/gpytorch/kernels/polynomial_kernel.py
@@ -93,7 +93,7 @@ def forward(
         if diag:
             return ((x1 * x2).sum(dim=-1) + self.offset).pow(self.power)
 
-        if x1.dim() == 2 and x2.dim() == 2:
+        if (x1.dim() == 2 and x2.dim() == 2) and offset.dim() == 2:
             return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)
         else:
             return (torch.matmul(x1, x2.transpose(-2, -1)) + offset).pow(self.power)


Test patch: