Repository: cornellius-gp/gpytorch
Issue: [Bug] PiecewisePolynomialKernel fails to put all tensors on the same GPU device

# 🐛 Bug

I was experimenting with the [tutorial of Exact GP multiple GPUs here](https://docs.gpytorch.ai/en/stable/examples/02_Scalable_Exact_GPs/Simple_MultiGPU_GP_Regression.html). However, when the base kernel was changed from **RBF kernel** to **piecewise polynomial kernel**, an error showed up that **tensors are not on the same device**.

## To reproduce

** Code snippet to reproduce **
```python
import torch
import gpytorch
from LBFGS import FullBatchLBFGS

import os
import numpy as np
import urllib.request
from scipy.io import loadmat
dataset = 'protein'
if not os.path.isfile(f'../../datasets/UCI/{dataset}.mat'):
    print(f'Downloading \\'{dataset}\\' UCI dataset...')
    urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1nRb8e7qooozXkNghC5eQS0JeywSXGX2S',
                               f'../../datasets/UCI/{dataset}.mat')

data = torch.Tensor(loadmat(f'../../datasets/UCI/{dataset}.mat')['data'])

n_train = 4000
train_x, train_y = data[:n_train, :-1], data[:n_train, -1]

n_devices = torch.cuda.device_count()
output_device = torch.device('cuda:0')
train_x, train_y = train_x.contiguous().to(output_device), train_y.contiguous().to(output_device)

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, n_devices):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        # change kernel here ----------------------------------------------------|
        base_covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.PiecewisePolynomialKernel())

        self.covar_module = gpytorch.kernels.MultiDeviceKernel(
            base_covar_module, device_ids=range(n_devices),
            output_device=output_device
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

def train(train_x,
          train_y,
          n_devices,
          output_device,
          checkpoint_size,
          preconditioner_size,
):
    likelihood = gpytorch.likelihoods.GaussianLikelihood().to(output_device)
    model = ExactGPModel(train_x, train_y, likelihood, n_devices).to(output_device)
    model.train()
    likelihood.train()

    optimizer = FullBatchLBFGS(model.parameters(), lr=0.1)
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

    with gpytorch.beta_features.checkpoint_kernel(checkpoint_size), \\
         gpytorch.settings.max_preconditioner_size(preconditioner_size):

        def closure():
            optimizer.zero_grad()
            output = model(train_x)
            loss = -mll(output, train_y)
            return loss

        loss = closure()
        loss.backward()

        options = {'closure': closure, 'current_loss': loss, 'max_ls': 10}
        loss, _, _, _, _, _, _, fail = optimizer.step(options)

        print(loss.item())
        
    return model, likelihood

_, _ = train(train_x, train_y,
             n_devices=n_devices, output_device=output_device,
             checkpoint_size=0, preconditioner_size=100)
```

** Stack trace/error message **
```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [4], in <cell line: 1>()
----> 1 _, _ = train(train_x, train_y,
      2              n_devices=n_devices, output_device=output_device,
      3              checkpoint_size=0, preconditioner_size=100)

Input In [3], in train(train_x, train_y, n_devices, output_device, checkpoint_size, preconditioner_size)
     41     return loss
     43 loss = closure()
---> 44 loss.backward()
     46 options = {'closure': closure, 'current_loss': loss, 'max_ls': 10}
     47 loss, _, _, _, _, _, _, fail = optimizer.step(options)

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/torch/_tensor.py:396, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
    387 if has_torch_function_unary(self):
    388     return handle_torch_function(
    389         Tensor.backward,
    390         (self,),
   (...)
    394         create_graph=create_graph,
    395         inputs=inputs)
--> 396 torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/torch/autograd/__init__.py:173, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
    168     retain_graph = create_graph
    170 # The reason we repeat same the comment below is that
    171 # some Python versions print out the first line of a multi-line function
    172 # calls in the traceback and some print out the last line
--> 173 Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
    174     tensors, grad_tensors_, retain_graph, create_graph, inputs,
    175     allow_unreachable=True, accumulate_grad=True)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cpu!
```


## System information

**Please complete the following information:**
- GPyTorch Version: 1.9.0
- PyTorch Version: 1.12.1
- Computer OS: Ubuntu 16.04.5 LTS
- CUDA version: 11.3
- CUDA devices: two NVIDIA 3090

## Additional context
I further experimented with training size and similar issue showed up when `n_train = 100` using **RBF kernel**. Please see the error message below. 

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [4], in <cell line: 1>()
----> 1 _, _ = train(train_x, train_y,
      2              n_devices=n_devices, output_device=output_device,
      3              checkpoint_size=0, preconditioner_size=100)

Input In [3], in train(train_x, train_y, n_devices, output_device, checkpoint_size, preconditioner_size)
     40     loss = -mll(output, train_y)
     41     return loss
---> 43 loss = closure()
     44 loss.backward()
     46 options = {'closure': closure, 'current_loss': loss, 'max_ls': 10}

Input In [3], in train.<locals>.closure()
     38 optimizer.zero_grad()
     39 output = model(train_x)
---> 40 loss = -mll(output, train_y)
     41 return loss

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/gpytorch/module.py:30, in Module.__call__(self, *inputs, **kwargs)
     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]

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/gpytorch/mlls/exact_marginal_log_likelihood.py:64, in ExactMarginalLogLikelihood.forward(self, function_dist, target, *params)
     62 # Get the log prob of the marginal distribution
     63 output = self.likelihood(function_dist, *params)
---> 64 res = output.log_prob(target)
     65 res = self._add_other_terms(res, params)
     67 # Scale by the amount of data we have

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/gpytorch/distributions/multivariate_normal.py:169, in MultivariateNormal.log_prob(self, value)
    167 # Get log determininant and first part of quadratic form
    168 covar = covar.evaluate_kernel()
--> 169 inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    171 res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])
    172 return res

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/operators/_linear_operator.py:1594, in LinearOperator.inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
   1592             will_need_cholesky = False
   1593     if will_need_cholesky:
-> 1594         cholesky = CholLinearOperator(TriangularLinearOperator(self.cholesky()))
   1595     return cholesky.inv_quad_logdet(
   1596         inv_quad_rhs=inv_quad_rhs,
   1597         logdet=logdet,
   1598         reduce_inv_quad=reduce_inv_quad,
   1599     )
   1601 # Short circuit to inv_quad function if we're not computing logdet

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/operators/_linear_operator.py:1229, in LinearOperator.cholesky(self, upper)
   1221 @_implements(torch.linalg.cholesky)
   1222 def cholesky(self, upper: bool = False) -> "TriangularLinearOperator":  # noqa F811
   1223     """
   1224     Cholesky-factorizes the LinearOperator.
   1225 
   1226     :param upper: Upper triangular or lower triangular factor (default: False).
   1227     :return: Cholesky factor (lower or upper triangular)
   1228     """
-> 1229     chol = self._cholesky(upper=False)
   1230     if upper:
   1231         chol = chol._transpose_nonbatch()

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/utils/memoize.py:59, in _cached.<locals>.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)

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/operators/_linear_operator.py:483, in LinearOperator._cholesky(self, upper)
    480 if any(isinstance(sub_mat, KeOpsLinearOperator) for sub_mat in evaluated_kern_mat._args):
    481     raise RuntimeError("Cannot run Cholesky with KeOps: it will either be really slow or not work.")
--> 483 evaluated_mat = evaluated_kern_mat.to_dense()
    485 # if the tensor is a scalar, we can just take the square root
    486 if evaluated_mat.size(-1) == 1:

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/utils/memoize.py:59, in _cached.<locals>.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)

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/operators/sum_linear_operator.py:68, in SumLinearOperator.to_dense(self)
     66 @cached
     67 def to_dense(self):
---> 68     return (sum(linear_op.to_dense() for linear_op in self.linear_ops)).contiguous()

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/operators/sum_linear_operator.py:68, in <genexpr>(.0)
     66 @cached
     67 def to_dense(self):
---> 68     return (sum(linear_op.to_dense() for linear_op in self.linear_ops)).contiguous()

File ~/anaconda3/envs/pyg/lib/python3.8/site-packages/linear_operator/operators/cat_linear_operator.py:378, in CatLinearOperator.to_dense(self)
    377 def to_dense(self):
--> 378     return torch.cat([to_dense(L) for L in self.linear_ops], dim=self.cat_dim)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument tensors in method wrapper_cat)
```


Changed files:
gpytorch/kernels/piecewise_polynomial_kernel.py

diff --git a/gpytorch/kernels/piecewise_polynomial_kernel.py b/gpytorch/kernels/piecewise_polynomial_kernel.py
index f6c9d5c5..8135979f 100644
--- a/gpytorch/kernels/piecewise_polynomial_kernel.py
+++ b/gpytorch/kernels/piecewise_polynomial_kernel.py
@@ -1,10 +1,32 @@
+import math
 from typing import Optional
 
 import torch
+from torch import Tensor
 
 from .kernel import Kernel
 
 
+def _fmax(r: Tensor, j: int, q: int) -> Tensor:
+    return torch.max(torch.tensor(0.0, dtype=r.dtype, device=r.device), 1 - r).pow(j + q)
+
+
+def _get_cov(r: Tensor, j: int, q: int) -> Tensor:
+    if q == 0:
+        return 1
+    if q == 1:
+        return (j + 1) * r + 1
+    if q == 2:
+        return 1 + (j + 2) * r + ((j + 4 * j + 3) / 3.0) * (r**2)
+    if q == 3:
+        return (
+            1
+            + (j + 3) * r
+            + ((6 * j**2 + 36 * j + 45) / 15.0) * r.square()
+            + ((j**3 + 9 * j**2 + 23 * j + 15) / 15.0) * (r**3)
+        )
+
+
 class PiecewisePolynomialKernel(Kernel):
     r"""
     Computes a covariance matrix based on the Piecewise Polynomial kernel
@@ -79,32 +101,14 @@ class PiecewisePolynomialKernel(Kernel):
             raise ValueError("q expected to be 0, 1, 2 or 3")
         self.q = q
 
-    def fmax(self, r, j, q):
-        return torch.max(torch.tensor(0.0), 1 - r).pow(j + q)
-
-    def get_cov(self, r, j, q):
-        if q == 0:
-            return 1
-        if q == 1:
-            return (j + 1) * r + 1
-        if q == 2:
-            return 1 + (j + 2) * r + ((j**2 + 4 * j + 3) / 3.0) * r**2
-        if q == 3:
-            return (
-                1
-                + (j + 3) * r
-                + ((6 * j**2 + 36 * j + 45) / 15.0) * r**2
-                + ((j**3 + 9 * j**2 + 23 * j + 15) / 15.0) * r**3
-            )
-
-    def forward(self, x1, x2, last_dim_is_batch=False, diag=False, **params):
+    def forward(self, x1: Tensor, x2: Tensor, last_dim_is_batch: bool = False, diag: bool = False, **params) -> Tensor:
         x1_ = x1.div(self.lengthscale)
         x2_ = x2.div(self.lengthscale)
         if last_dim_is_batch is True:
             D = x1.shape[1]
         else:
             D = x1.shape[-1]
-        j = torch.floor(torch.tensor(D / 2.0)) + self.q + 1
+        j = math.floor(D / 2.0) + self.q + 1
         if last_dim_is_batch and diag:
             r = self.covar_dist(x1_, x2_, last_dim_is_batch=True, diag=True)
         elif diag:
@@ -113,5 +117,5 @@ class PiecewisePolynomialKernel(Kernel):
             r = self.covar_dist(x1_, x2_, last_dim_is_batch=True)
         else:
             r = self.covar_dist(x1_, x2_)
-        cov_matrix = self.fmax(r, j, self.q) * self.get_cov(r, j, self.q)
+        cov_matrix = _fmax(r, j, self.q) * _get_cov(r, j, self.q)
         return cov_matrix