Repository: cornellius-gp/gpytorch
Issue: [Bug] torch.float62 raises error in `GridInterpolationKernel`

# 🐛 Bug

I want to sample from the prior distribution with precision torch.float62 
However, during sampling with KISS-GP a dtype error is raised
if I manually design a kernel (very similar to the RBF Kernel)
that is included in the `GridInterpolationKernel`.

Changing the test data size from `2500x2` to `100x2`, no error will occur.
```python
x = torch.meshgrid(
    torch.linspace(0, 10 - 1, 10) * 1.,
    torch.linspace(0, 10 - 1, 10) * 1.,
    indexing="xy",
)
x = torch.cat(
    (
        x[0].contiguous().view(x[0].numel(), 1),
        x[1].contiguous().view(x[1].numel(), 1),
    ),
    dim=1,
)
``` 
## To reproduce

```python
import torch
import gpytorch


torch.set_default_dtype(torch.float64)


def postprocess_rot(dist_mat):
    return dist_mat.mul_(-1.0).exp_()

class TestKernel(gpytorch.kernels.Kernel):

    is_stationary = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def forward(self, x1, x2, **params):
        x1_ = x1.div_(torch.tensor([10., 1.]))
        x2_ = x2.div_(torch.tensor([10., 1.]))
        return self.covar_dist(
            x1_, x2_, square_dist=False, dist_postprocess_func=postprocess_rot, **params
        )

class ExactGP(gpytorch.models.ExactGP):

    def __init__(self, **kwargs):
        super().__init__(None, None, gpytorch.likelihoods.GaussianLikelihood())
        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.GridInterpolationKernel(
            TestKernel(ard_num_dims=2, **kwargs),
            grid_size=100,
            num_dims=2
            )

    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 = torch.meshgrid(
    torch.linspace(0, 50 - 1, 50) * 1.,
    torch.linspace(0, 50 - 1, 50) * 1.,
    indexing="xy",
)
x = torch.cat(
    (
        x[0].contiguous().view(x[0].numel(), 1),
        x[1].contiguous().view(x[1].numel(), 1),
    ),
    dim=1,
)

model = ExactGP()
model.eval()

with torch.no_grad(), gpytorch.settings.fast_pred_var(), gpytorch.settings.max_root_decomposition_size(100):
    with gpytorch.settings.fast_pred_samples():
        samples = model(x).rsample(torch.Size([1]))

```

** Error message **
```
expected scalar type Double but found Float
```

## System information

torch=1.13.0
gpytorch=1.9.0

Changed files:
gpytorch/likelihoods/gaussian_likelihood.py, test/examples/test_kissgp_gp_regression.py, test/examples/test_kissgp_kronecker_product_regression.py, gpytorch/kernels/newton_girard_additive_kernel.py, gpytorch/utils/grid.py

diff --git a/gpytorch/kernels/newton_girard_additive_kernel.py b/gpytorch/kernels/newton_girard_additive_kernel.py
index 13a9e5ae..d01c244e 100644
--- a/gpytorch/kernels/newton_girard_additive_kernel.py
+++ b/gpytorch/kernels/newton_girard_additive_kernel.py
@@ -90,14 +90,14 @@ class NewtonGirardAdditiveKernel(Kernel):
         s_k = kern_values.unsqueeze(kernel_dim - 1).pow(kvals).sum(dim=kernel_dim)
 
         # just the constant -1
-        m1 = torch.tensor([-1], dtype=torch.float, device=kern_values.device)
+        m1 = torch.tensor([-1], dtype=x1.dtype, device=kern_values.device)
 
         shape = [1 for _ in range(len(kern_values.shape))]
         shape[kernel_dim] = -1
         for deg in range(1, self.max_degree + 1):  # deg goes from 1 to R (it's 1-indexed!)
             # we avg over k [1, ..., deg] (-1)^(k-1)e_{deg-k} s_{k}
 
-            ks = torch.arange(1, deg + 1, device=kern_values.device, dtype=torch.float).reshape(*shape)  # use for pow
+            ks = torch.arange(1, deg + 1, device=kern_values.device, dtype=x1.dtype).reshape(*shape)  # use for pow
             kslong = torch.arange(1, deg + 1, device=kern_values.device, dtype=torch.long)  # use for indexing
 
             # note that s_k is 0-indexed, so we must subtract 1 from kslong
diff --git a/gpytorch/likelihoods/gaussian_likelihood.py b/gpytorch/likelihoods/gaussian_likelihood.py
index 318bd655..b366fa2f 100644
--- a/gpytorch/likelihoods/gaussian_likelihood.py
+++ b/gpytorch/likelihoods/gaussian_likelihood.py
@@ -308,7 +308,10 @@ class DirichletClassificationLikelihood(FixedNoiseGaussianLikelihood):
         >>> pred_y = likelihood(gp_model(test_x), targets=labels)
     """
 
-    def _prepare_targets(self, targets, alpha_epsilon=0.01, dtype=torch.float):
+    def _prepare_targets(self, targets: torch.Tensor, alpha_epsilon: float = 0.01, dtype: Optional[torch.dtype] = None) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        if dtype is None:
+            dtype = torch.get_default_dtype()
+
         num_classes = int(targets.max() + 1)
         # set alpha = \\alpha_\\epsilon
         alpha = alpha_epsilon * torch.ones(targets.shape[-1], num_classes, device=targets.device, dtype=dtype)
diff --git a/gpytorch/utils/grid.py b/gpytorch/utils/grid.py
index 7cd57d87..12c76e3a 100644
--- a/gpytorch/utils/grid.py
+++ b/gpytorch/utils/grid.py
@@ -2,7 +2,7 @@
 
 import math
 import warnings
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 import torch
 
@@ -130,8 +130,8 @@ def create_grid(
     grid_sizes: List[int],
     grid_bounds: List[Tuple[float, float]],
     extend: bool = True,
-    device="cpu",
-    dtype=torch.float,
+    device: str = "cpu",
+    dtype: Optional[torch.dtype] = None,
 ) -> List[torch.Tensor]:
     """
     Creates a grid represented by a list of 1D Tensors representing the
@@ -141,16 +141,14 @@ def create_grid(
     which can be important for getting good grid interpolations.
 
     :param grid_sizes: Sizes of each grid dimension
-    :type grid_sizes: List[int]
     :param grid_bounds: Lower and upper bounds of each grid dimension
-    :type grid_sizes: List[Tuple[float, float]]
     :param device: target device for output (default: cpu)
-    :type device: torch.device, optional
     :param dtype: target dtype for output (default: torch.float)
-    :type dtype: torch.dtype, optional
     :return: Grid points for each dimension. Grid points are stored in a :obj:`torch.Tensor` with shape `grid_sizes[i]`.
     :rtype: List[torch.Tensor]
     """
+    if dtype is None:
+        dtype = torch.get_default_dtype()
     grid = []
     for i in range(len(grid_bounds)):
         grid_diff = float(grid_bounds[i][1] - grid_bounds[i][0]) / (grid_sizes[i] - 2)
diff --git a/test/examples/test_kissgp_gp_regression.py b/test/examples/test_kissgp_gp_regression.py
index 6708d87a..762530be 100644
--- a/test/examples/test_kissgp_gp_regression.py
+++ b/test/examples/test_kissgp_gp_regression.py
@@ -1,7 +1,5 @@
 #!/usr/bin/env python3
 
-import os
-import random
 import unittest
 from math import exp, pi
 
@@ -12,6 +10,7 @@ from gpytorch.kernels import GridInterpolationKernel, RBFKernel, ScaleKernel
 from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.means import ConstantMean
 from gpytorch.priors import SmoothedBoxPrior
+from gpytorch.test.base_test_case import BaseTestCase
 from gpytorch.test.utils import least_used_cuda_device
 from torch import optim
 
@@ -44,18 +43,8 @@ class GPRegressionModel(gpytorch.models.ExactGP):
         return MultivariateNormal(mean_x, covar_x)
 
 
-class TestKISSGPRegression(unittest.TestCase):
-    def setUp(self):
-        if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(0)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(0)
-            random.seed(0)
-
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
+class TestKISSGPRegression(unittest.TestCase, BaseTestCase):
+    seed = 0
 
     def test_kissgp_gp_mean_abs_error(self):
         train_x, train_y, test_x, test_y = make_data()
diff --git a/test/examples/test_kissgp_kronecker_product_regression.py b/test/examples/test_kissgp_kronecker_product_regression.py
index 36457ff4..61d13ffd 100644
--- a/test/examples/test_kissgp_kronecker_product_regression.py
+++ b/test/examples/test_kissgp_kronecker_product_regression.py
@@ -2,8 +2,6 @@
 
 from math import pi
 
-import os
-import random
 import torch
 import unittest
 
@@ -13,6 +11,7 @@ from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.means import ConstantMean
 from gpytorch.priors import SmoothedBoxPrior
 from gpytorch.distributions import MultivariateNormal
+from gpytorch.test.base_test_case import BaseTestCase
 from torch import optim
 
 # Simple training data: let's try to learn a sine function,
@@ -52,18 +51,8 @@ class GPRegressionModel(gpytorch.models.ExactGP):
         return MultivariateNormal(mean_x, covar_x)
 
 
-class TestKISSGPKroneckerProductRegression(unittest.TestCase):
-    def setUp(self):
-        if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(0)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(0)
-            random.seed(0)
-
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
+class TestKISSGPKroneckerProductRegression(unittest.TestCase, BaseTestCase):
+    seed = 0
 
     def test_kissgp_gp_mean_abs_error(self):
         likelihood = GaussianLikelihood()
@@ -99,6 +88,22 @@ class TestKISSGPKroneckerProductRegression(unittest.TestCase):
             mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
             self.assertLess(mean_abs_error.squeeze().item(), 0.2)
 
+            # Try drawing a sample - make sure there's no errors
+            with torch.no_grad(), gpytorch.settings.max_root_decomposition_size(100):
+                with gpytorch.settings.fast_pred_samples():
+                    gp_model(train_x).rsample(torch.Size([1]))
+
+
+class TestKISSGPKroneckerProductRegressionDouble(TestKISSGPKroneckerProductRegression):
+    def setUp(self):
+        super().setUp
+        self.default_dtype = torch.get_default_dtype()
+        torch.set_default_dtype(torch.float64)
+
+    def tearDown(self):
+        super().tearDown()
+        torch.set_default_dtype(self.default_dtype)
+
 
 if __name__ == "__main__":
     unittest.main()