We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `gpytorch/kernels/polynomial_kernel.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 from typing import Optional
4 
5 import torch
6 
7 from ..constraints import Interval, Positive
8 from ..priors import Prior
9 from .kernel import Kernel
10 
11 
12 class PolynomialKernel(Kernel):
13     r"""
14     Computes a covariance matrix based on the Polynomial kernel
15     between inputs :math:`\mathbf{x_1}` and :math:`\mathbf{x_2}`:
16 
17     .. math::
18         \begin{equation*}
19             k_\text{Poly}(\mathbf{x_1}, \mathbf{x_2}) = (\mathbf{x_1}^\top
20             \mathbf{x_2} + c)^{d}.
21         \end{equation*}
22 
23     where
24 
25     * :math:`c` is an :attr:`offset` parameter.
26 
27     Args:
28         :attr:`offset_prior` (:class:`gpytorch.priors.Prior`):
29             Prior over the offset parameter (default `None`).
30         :attr:`offset_constraint` (Constraint, optional):
31             Constraint to place on offset parameter. Default: `Positive`.
32         :attr:`active_dims` (list):
33             List of data dimensions to operate on.
34             `len(active_dims)` should equal `num_dimensions`.
35     """
36 
37     def __init__(
38         self,
39         power: int,
40         offset_prior: Optional[Prior] = None,
41         offset_constraint: Optional[Interval] = None,
42         **kwargs,
43     ):
44         super().__init__(**kwargs)
45         if offset_constraint is None:
46             offset_constraint = Positive()
47 
48         self.register_parameter(name="raw_offset", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)))
49 
50         # We want the power to be a float so we dont have to worry about its device / dtype.
51         if torch.is_tensor(power):
52             if power.numel() > 1:
53                 raise RuntimeError("Cant create a Polynomial kernel with more than one power")
54             else:
55                 power = power.item()
56 
57         self.power = power
58 
59         if offset_prior is not None:
60             if not isinstance(offset_prior, Prior):
61                 raise TypeError("Expected gpytorch.priors.Prior but got " + type(offset_prior).__name__)
62             self.register_prior("offset_prior", offset_prior, lambda m: m.offset, lambda m, v: m._set_offset(v))
63 
64         self.register_constraint("raw_offset", offset_constraint)
65 
66     @property
67     def offset(self) -> torch.Tensor:
68         return self.raw_offset_constraint.transform(self.raw_offset)
69 
70     @offset.setter
71     def offset(self, value: torch.Tensor) -> None:
72         self._set_offset(value)
73 
74     def _set_offset(self, value: torch.Tensor) -> None:
75         if not torch.is_tensor(value):
76             value = torch.as_tensor(value).to(self.raw_offset)
77         self.initialize(raw_offset=self.raw_offset_constraint.inverse_transform(value))
78 
79     def forward(
80         self,
81         x1: torch.Tensor,
82         x2: torch.Tensor,
83         diag: Optional[bool] = False,
84         last_dim_is_batch: Optional[bool] = False,
85         **params,
86     ) -> torch.Tensor:
87         offset = self.offset.view(*self.batch_shape, 1, 1)
88 
89         if last_dim_is_batch:
90             x1 = x1.transpose(-1, -2).unsqueeze(-1)
91             x2 = x2.transpose(-1, -2).unsqueeze(-1)
92 
93         if diag:
94             return ((x1 * x2).sum(dim=-1) + self.offset).pow(self.power)
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)
100
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
         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)

{"golden_diff": "diff --git a/gpytorch/kernels/polynomial_kernel.py b/gpytorch/kernels/polynomial_kernel.py\n--- a/gpytorch/kernels/polynomial_kernel.py\n+++ b/gpytorch/kernels/polynomial_kernel.py\n@@ -93,7 +93,7 @@\n         if diag:\n             return ((x1 * x2).sum(dim=-1) + self.offset).pow(self.power)\n \n-        if x1.dim() == 2 and x2.dim() == 2:\n+        if (x1.dim() == 2 and x2.dim() == 2) and offset.dim() == 2:\n             return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)\n         else:\n             return (torch.matmul(x1, x2.transpose(-2, -1)) + offset).pow(self.power)\n", "issue": "[Bug] IF-construction leads to RuntimeError in polynomial kernel \n# \ud83d\udc1b Bug\r\n\r\nI 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.\r\nIt crushes with RuntimeError while inference. And from what I see in the code there might be several reasons why:\r\n- misprint (if and else should be vice versa)\r\n- pytorch changed the behavior of `torch.addmm`\r\n- or maybe there is smth I do not understand :)\r\n\r\nPlease, look at the text below.\r\n\r\n## To reproduce\r\n\r\n** Code snippet to reproduce **\r\n```python\r\nimport torch\r\nimport gpytorch\r\n\r\nclass MultitaskGPModel(gpytorch.models.ApproximateGP):\r\n    def __init__(self, grid_size=4, num_latents=2, grid_bounds=[(0., 31.)], num_tasks=256):\r\n        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(\r\n            grid_size, batch_shape=torch.Size([num_latents])\r\n        )\r\n        variational_strategy = gpytorch.variational.LMCVariationalStrategy(\r\n            gpytorch.variational.GridInterpolationVariationalStrategy(self, grid_size, grid_bounds, variational_distribution),\r\n            num_tasks=num_tasks,\r\n            num_latents=num_latents,\r\n            latent_dim=-1\r\n        )\r\n        super().__init__(variational_strategy)\r\n\r\n        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([num_latents]))\r\n        self.covar_module = gpytorch.kernels.ScaleKernel(\r\n            gpytorch.kernels.PolynomialKernel(batch_shape=torch.Size([num_latents]), power=2),\r\n            batch_shape=torch.Size([num_latents])\r\n        )\r\n        \r\n    def forward(self, x):\r\n        mean_x = self.mean_module(x)\r\n        covar_x = self.covar_module(x)\r\n        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\r\n    \r\ngp_layer = MultitaskGPModel()\r\nlikelihood = gpytorch.likelihoods.MultitaskGaussianLikelihood(num_tasks=256)\r\n\r\nx = torch.randint(32, size=(1,), dtype=torch.float32).unsqueeze(1)\r\n\r\nnoise = gp_layer(x)\r\nnoise = likelihood(noise)\r\nnoise = noise.rsample((2, 2))\r\n```\r\n\r\n** Stack trace/error message **\r\n```\r\n---------------------------------------------------------------------------\r\nRuntimeError                              Traceback (most recent call last)\r\n/tmp/ipykernel_15909/3264522926.py in <module>\r\n     31 x = torch.randint(32, size=(1,), dtype=torch.float32).unsqueeze(1)\r\n     32 \r\n---> 33 noise = gp_layer(x)\r\n     34 noise = likelihood(noise)\r\n     35 noise = noise.rsample((2, 2))\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/models/approximate_gp.py in __call__(self, inputs, prior, **kwargs)\r\n     79         if inputs.dim() == 1:\r\n     80             inputs = inputs.unsqueeze(-1)\r\n---> 81         return self.variational_strategy(inputs, prior=prior, **kwargs)\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/variational/lmc_variational_strategy.py in __call__(self, x, task_indices, prior, **kwargs)\r\n    186             or ~gpytorch.distributions.MultivariateNormal (... x N)\r\n    187         \"\"\"\r\n--> 188         latent_dist = self.base_variational_strategy(x, prior=prior, **kwargs)\r\n    189         num_batch = len(latent_dist.batch_shape)\r\n    190         latent_dim = num_batch + self.latent_dim\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/variational/_variational_strategy.py in __call__(self, x, prior, **kwargs)\r\n    108         # (Maybe) initialize variational distribution\r\n    109         if not self.variational_params_initialized.item():\r\n--> 110             prior_dist = self.prior_distribution\r\n    111             self._variational_distribution.initialize_variational_distribution(prior_dist)\r\n    112             self.variational_params_initialized.fill_(1)\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)\r\n     57         kwargs_pkl = pickle.dumps(kwargs)\r\n     58         if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):\r\n---> 59             return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)\r\n     60         return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)\r\n     61 \r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/variational/grid_interpolation_variational_strategy.py in prior_distribution(self)\r\n     74     def prior_distribution(self):\r\n     75         out = self.model.forward(self.inducing_points)\r\n---> 76         res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter())\r\n     77         return res\r\n     78 \r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in add_jitter(self, jitter_val)\r\n    275 \r\n    276     def add_jitter(self, jitter_val=1e-3):\r\n--> 277         return self.evaluate_kernel().add_jitter(jitter_val)\r\n    278 \r\n    279     def _unsqueeze_batch(self, dim):\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)\r\n     57         kwargs_pkl = pickle.dumps(kwargs)\r\n     58         if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):\r\n---> 59             return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)\r\n     60         return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)\r\n     61 \r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)\r\n    330             temp_active_dims = self.kernel.active_dims\r\n    331             self.kernel.active_dims = None\r\n--> 332             res = self.kernel(\r\n    333                 x1,\r\n    334                 x2,\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/kernels/kernel.py in __call__(self, x1, x2, diag, last_dim_is_batch, **params)\r\n    400                 res = LazyEvaluatedKernelTensor(x1_, x2_, kernel=self, last_dim_is_batch=last_dim_is_batch, **params)\r\n    401             else:\r\n--> 402                 res = lazify(super(Kernel, self).__call__(x1_, x2_, last_dim_is_batch=last_dim_is_batch, **params))\r\n    403             return res\r\n    404 \r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)\r\n     28 \r\n     29     def __call__(self, *inputs, **kwargs):\r\n---> 30         outputs = self.forward(*inputs, **kwargs)\r\n     31         if isinstance(outputs, list):\r\n     32             return [_validate_module_outputs(output) for output in outputs]\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/kernels/scale_kernel.py in forward(self, x1, x2, last_dim_is_batch, diag, **params)\r\n    101 \r\n    102     def forward(self, x1, x2, last_dim_is_batch=False, diag=False, **params):\r\n--> 103         orig_output = self.base_kernel.forward(x1, x2, diag=diag, last_dim_is_batch=last_dim_is_batch, **params)\r\n    104         outputscales = self.outputscale\r\n    105         if last_dim_is_batch:\r\n\r\n~/anaconda3/lib/python3.8/site-packages/gpytorch/kernels/polynomial_kernel.py in forward(self, x1, x2, diag, last_dim_is_batch, **params)\r\n     95 \r\n     96         if x1.dim() == 2 and x2.dim() == 2:\r\n---> 97             return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)\r\n     98         else:\r\n     99             return (torch.matmul(x1, x2.transpose(-2, -1)) + offset).pow(self.power)\r\n\r\nRuntimeError: 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)\r\n```\r\n\r\n## System information\r\n\r\nGPyTorch Version: 1.6.0\r\nPyTorch Version : 1.11.0+cu113\r\nComputer OS: Ubuntu 18.04.6 LTS \r\n\r\n## Additional context\r\nI see in [`polynomial_kernel.py`](https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/kernels/polynomial_kernel.py) this piece of code:\r\n\r\n```python\r\noffset = self.offset.view(*self.batch_shape, 1, 1)\r\n...\r\nif x1.dim() == 2 and x2.dim() == 2:\r\n    return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)\r\n```\r\n\r\nSo 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.\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nfrom typing import Optional\n\nimport torch\n\nfrom ..constraints import Interval, Positive\nfrom ..priors import Prior\nfrom .kernel import Kernel\n\n\nclass PolynomialKernel(Kernel):\n    r\"\"\"\n    Computes a covariance matrix based on the Polynomial kernel\n    between inputs :math:`\\mathbf{x_1}` and :math:`\\mathbf{x_2}`:\n\n    .. math::\n        \\begin{equation*}\n            k_\\text{Poly}(\\mathbf{x_1}, \\mathbf{x_2}) = (\\mathbf{x_1}^\\top\n            \\mathbf{x_2} + c)^{d}.\n        \\end{equation*}\n\n    where\n\n    * :math:`c` is an :attr:`offset` parameter.\n\n    Args:\n        :attr:`offset_prior` (:class:`gpytorch.priors.Prior`):\n            Prior over the offset parameter (default `None`).\n        :attr:`offset_constraint` (Constraint, optional):\n            Constraint to place on offset parameter. Default: `Positive`.\n        :attr:`active_dims` (list):\n            List of data dimensions to operate on.\n            `len(active_dims)` should equal `num_dimensions`.\n    \"\"\"\n\n    def __init__(\n        self,\n        power: int,\n        offset_prior: Optional[Prior] = None,\n        offset_constraint: Optional[Interval] = None,\n        **kwargs,\n    ):\n        super().__init__(**kwargs)\n        if offset_constraint is None:\n            offset_constraint = Positive()\n\n        self.register_parameter(name=\"raw_offset\", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)))\n\n        # We want the power to be a float so we dont have to worry about its device / dtype.\n        if torch.is_tensor(power):\n            if power.numel() > 1:\n                raise RuntimeError(\"Cant create a Polynomial kernel with more than one power\")\n            else:\n                power = power.item()\n\n        self.power = power\n\n        if offset_prior is not None:\n            if not isinstance(offset_prior, Prior):\n                raise TypeError(\"Expected gpytorch.priors.Prior but got \" + type(offset_prior).__name__)\n            self.register_prior(\"offset_prior\", offset_prior, lambda m: m.offset, lambda m, v: m._set_offset(v))\n\n        self.register_constraint(\"raw_offset\", offset_constraint)\n\n    @property\n    def offset(self) -> torch.Tensor:\n        return self.raw_offset_constraint.transform(self.raw_offset)\n\n    @offset.setter\n    def offset(self, value: torch.Tensor) -> None:\n        self._set_offset(value)\n\n    def _set_offset(self, value: torch.Tensor) -> None:\n        if not torch.is_tensor(value):\n            value = torch.as_tensor(value).to(self.raw_offset)\n        self.initialize(raw_offset=self.raw_offset_constraint.inverse_transform(value))\n\n    def forward(\n        self,\n        x1: torch.Tensor,\n        x2: torch.Tensor,\n        diag: Optional[bool] = False,\n        last_dim_is_batch: Optional[bool] = False,\n        **params,\n    ) -> torch.Tensor:\n        offset = self.offset.view(*self.batch_shape, 1, 1)\n\n        if last_dim_is_batch:\n            x1 = x1.transpose(-1, -2).unsqueeze(-1)\n            x2 = x2.transpose(-1, -2).unsqueeze(-1)\n\n        if diag:\n            return ((x1 * x2).sum(dim=-1) + self.offset).pow(self.power)\n\n        if x1.dim() == 2 and x2.dim() == 2:\n            return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)\n        else:\n            return (torch.matmul(x1, x2.transpose(-2, -1)) + offset).pow(self.power)\n", "path": "gpytorch/kernels/polynomial_kernel.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nfrom typing import Optional\n\nimport torch\n\nfrom ..constraints import Interval, Positive\nfrom ..priors import Prior\nfrom .kernel import Kernel\n\n\nclass PolynomialKernel(Kernel):\n    r\"\"\"\n    Computes a covariance matrix based on the Polynomial kernel\n    between inputs :math:`\\mathbf{x_1}` and :math:`\\mathbf{x_2}`:\n\n    .. math::\n        \\begin{equation*}\n            k_\\text{Poly}(\\mathbf{x_1}, \\mathbf{x_2}) = (\\mathbf{x_1}^\\top\n            \\mathbf{x_2} + c)^{d}.\n        \\end{equation*}\n\n    where\n\n    * :math:`c` is an :attr:`offset` parameter.\n\n    Args:\n        :attr:`offset_prior` (:class:`gpytorch.priors.Prior`):\n            Prior over the offset parameter (default `None`).\n        :attr:`offset_constraint` (Constraint, optional):\n            Constraint to place on offset parameter. Default: `Positive`.\n        :attr:`active_dims` (list):\n            List of data dimensions to operate on.\n            `len(active_dims)` should equal `num_dimensions`.\n    \"\"\"\n\n    def __init__(\n        self,\n        power: int,\n        offset_prior: Optional[Prior] = None,\n        offset_constraint: Optional[Interval] = None,\n        **kwargs,\n    ):\n        super().__init__(**kwargs)\n        if offset_constraint is None:\n            offset_constraint = Positive()\n\n        self.register_parameter(name=\"raw_offset\", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1)))\n\n        # We want the power to be a float so we dont have to worry about its device / dtype.\n        if torch.is_tensor(power):\n            if power.numel() > 1:\n                raise RuntimeError(\"Cant create a Polynomial kernel with more than one power\")\n            else:\n                power = power.item()\n\n        self.power = power\n\n        if offset_prior is not None:\n            if not isinstance(offset_prior, Prior):\n                raise TypeError(\"Expected gpytorch.priors.Prior but got \" + type(offset_prior).__name__)\n            self.register_prior(\"offset_prior\", offset_prior, lambda m: m.offset, lambda m, v: m._set_offset(v))\n\n        self.register_constraint(\"raw_offset\", offset_constraint)\n\n    @property\n    def offset(self) -> torch.Tensor:\n        return self.raw_offset_constraint.transform(self.raw_offset)\n\n    @offset.setter\n    def offset(self, value: torch.Tensor) -> None:\n        self._set_offset(value)\n\n    def _set_offset(self, value: torch.Tensor) -> None:\n        if not torch.is_tensor(value):\n            value = torch.as_tensor(value).to(self.raw_offset)\n        self.initialize(raw_offset=self.raw_offset_constraint.inverse_transform(value))\n\n    def forward(\n        self,\n        x1: torch.Tensor,\n        x2: torch.Tensor,\n        diag: Optional[bool] = False,\n        last_dim_is_batch: Optional[bool] = False,\n        **params,\n    ) -> torch.Tensor:\n        offset = self.offset.view(*self.batch_shape, 1, 1)\n\n        if last_dim_is_batch:\n            x1 = x1.transpose(-1, -2).unsqueeze(-1)\n            x2 = x2.transpose(-1, -2).unsqueeze(-1)\n\n        if diag:\n            return ((x1 * x2).sum(dim=-1) + self.offset).pow(self.power)\n\n        if (x1.dim() == 2 and x2.dim() == 2) and offset.dim() == 2:\n            return torch.addmm(offset, x1, x2.transpose(-2, -1)).pow(self.power)\n        else:\n            return (torch.matmul(x1, x2.transpose(-2, -1)) + offset).pow(self.power)\n", "path": "gpytorch/kernels/polynomial_kernel.py"}]}