We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Bug] Gradient computation with KeOps kernel only works with low number of training data, fails otherwise
# 🐛 Bug

Hi,

First of all, I would like to extend my thanks to all the developers for all the efforts you have put into both the research and this great package. 

Consider a GP with a KeOps kernel (e.g. gpytorch.kernels.keops.RBFKernel). If I train it with N=100 number of points, then the gradient of predictive mean can be obtained by torch.autograd.grad or .backward(), but set N=500 and an error will be thrown that the input tensor was not used in the graph. I have tested the script on two separate machines and a colab instance. Using the GPyTorch standard kernels will not run into this issue. I spent a good deal of time pinpointing what was wrong from bigger chunks of code and this seemed to be the issue. The gradient link seems to cut off between the covariance output of the KeOps kernel and the input (covar.x1).  

I have provided a minimal code right below that should quickly give you an idea of this somewhat strange behavior. It contains two test cases with N=100 (which passes) and 500 (fails). The code is from GPyTorch regression examples, I have only added the GP kernel and a few lines for calculating gradients at the end. 

## To reproduce

```python
import math
import torch
import gpytorch
import time

# We will use the simplest form of GP model, exact inference with gpytorch.kernels.keops.RBFKernel
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.keops.RBFKernel())

    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_and_eval_GP(N = 100):
    """
    inputs:
    N (int): Number of training points
    """
    # make train/val/test
    # Training data is 100 points in [0,1] inclusive regularly spaced
    train_x = torch.linspace(0, 1, N)
    # True function is sin(2*pi*x) with Gaussian noise
    train_y = torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * math.sqrt(0.04)
    # normalize features
    mean = train_x.mean()
    std = train_x.std() + 1e-6 # prevent dividing by 0
    train_x = (train_x - mean) / std

    # normalize labels
    mean, std = train_y.mean(),train_y.std()
    train_y = (train_y - mean) / std

    # make continguous
    train_x, train_y = train_x.contiguous(), train_y.contiguous()

    output_device = torch.device('cuda:0')

    train_x, train_y = train_x.to(output_device), train_y.to(output_device)

    # initialize likelihood and model
    likelihood = gpytorch.likelihoods.GaussianLikelihood().to(output_device)
    model = ExactGPModel(train_x, train_y, likelihood).to(output_device)

    # Find optimal model hyperparameters
    model.train()
    likelihood.train()

    # Use the adam optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)  # Includes GaussianLikelihood parameters

    # "Loss" for GPs - the marginal log likelihood
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

    training_iter = 20
    for i in range(training_iter):
        # Zero gradients from previous iteration
        optimizer.zero_grad()
        # Output from model
        output = model(train_x)
        # Calc loss and backprop gradients
        loss = -mll(output, train_y)
        loss.backward()
        optimizer.step()
    print(f'GP model trained.')

    # Get into evaluation (predictive posterior) mode
    model.eval()
    likelihood.eval()

    # Test points are regularly spaced along [0,1]
    test_x = torch.linspace(0, 1, 51, requires_grad=True).to(output_device).contiguous()
    
    # Make predictions by feeding model through likelihood
    with gpytorch.settings.fast_pred_var():
        observed_pred = likelihood(model(test_x))
        assert torch.autograd.grad(observed_pred.mean.sum(), test_x, retain_graph=True) is not None
        print('gradient exists:')
        print(torch.autograd.grad(observed_pred.mean.sum(), test_x, retain_graph=True))

if __name__ == "__main__":
    Ns = [100, 500] #test cases
    for n in Ns:
        try:
            print(f'testing with {n} points...')
            train_and_eval_GP(N = n) 
            print('success!')
        except Exception as e:
            print('failed.')
            print(e)
```

** Stack trace/error message **
```
testing with 100 points...
GP model trained.
gradient exists:
(tensor([-2.6629e+00, -2.6507e+00, -2.6344e+00, -2.6138e+00, -2.5891e+00,
        ...,
         9.1844e-01], device='cuda:0'),)
success!
testing with 500 points...
GP model trained.
failed.
One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
```

## Expected Behavior

Setting the kernel to the standard, non-KeOps kernel (gpytorch.kernels.RBFKernel) we get the gradients for the second case, shown below. However, I can't simply use it since I'm working on a larger dataset that will run out of memory if I do so.

```
testing with 100 points...
GP model trained.
gradient exists:
(tensor([-2.5885, -2.5870, -2.5819, -2.5732, -2.5609, -2.5449, -2.5254, -2.5022,
        ...,
         0.6977,  0.7880,  0.8762], device='cuda:0'),)
success!
testing with 500 points...
GP model trained.
gradient exists:
(tensor([-2.4751, -2.4741, -2.4698, -2.4622, -2.4514, -2.4373, -2.4199, -2.3993,
        ...,
         0.4727,  0.5600,  0.6462], device='cuda:0'),)
success!
```
## System information

**Please complete the following information:**
- <!-- GPyTorch Version (run `print(gpytorch.__version__)` --> GPyTorch 1.6.0
- <!-- PyTorch Version (run `print(torch.__version__)` --> PyTorch  1.10.1
- <!-- KeOps Version --> KeOps 1.5
- <!-- Computer OS --> Ubuntu 20.04.3 LTS

## Additional context
I know that the issue might be unrelated to GPyTorch as this clearly stems from the keops kernel. It is however difficult to track, so I thought I'd report it here. 


--- 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/keops/rbf_kernel.py`
Content:
```
1 #!/usr/bin/env python3
2 import torch
3 
4 from ... import settings
5 from ...lazy import KeOpsLazyTensor
6 from ..rbf_kernel import postprocess_rbf
7 from .keops_kernel import KeOpsKernel
8 
9 try:
10     from pykeops.torch import LazyTensor as KEOLazyTensor
11 
12     class RBFKernel(KeOpsKernel):
13         """
14         Implements the RBF kernel using KeOps as a driver for kernel matrix multiplies.
15 
16         This class can be used as a drop in replacement for gpytorch.kernels.RBFKernel in most cases, and supports
17         the same arguments. There are currently a few limitations, for example a lack of batch mode support. However,
18         most other features like ARD will work.
19         """
20 
21         has_lengthscale = True
22 
23         def _nonkeops_covar_func(self, x1, x2, diag=False):
24             return self.covar_dist(
25                 x1, x2, square_dist=True, diag=diag, dist_postprocess_func=postprocess_rbf, postprocess=True
26             )
27 
28         def covar_func(self, x1, x2, diag=False):
29             # We only should use KeOps on big kernel matrices
30             # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)
31             # then don't apply KeOps
32             if (
33                 diag
34                 or x1.size(-2) < settings.max_cholesky_size.value()
35                 or x2.size(-2) < settings.max_cholesky_size.value()
36             ):
37                 return self._nonkeops_covar_func(x1, x2, diag=diag)
38 
39             with torch.autograd.enable_grad():
40                 x1_ = KEOLazyTensor(x1[..., :, None, :])
41                 x2_ = KEOLazyTensor(x2[..., None, :, :])
42 
43                 K = (-((x1_ - x2_) ** 2).sum(-1) / 2).exp()
44 
45                 return K
46 
47         def forward(self, x1, x2, diag=False, **params):
48             x1_ = x1.div(self.lengthscale)
49             x2_ = x2.div(self.lengthscale)
50 
51             covar_func = lambda x1, x2, diag=diag: self.covar_func(x1, x2, diag)
52 
53             if diag:
54                 return covar_func(x1_, x2_, diag=True)
55 
56             return KeOpsLazyTensor(x1_, x2_, covar_func)
57 
58 
59 except ImportError:
60 
61     class RBFKernel(KeOpsKernel):
62         def __init__(self, *args, **kwargs):
63             super().__init__()
64
```

Path: `gpytorch/kernels/keops/matern_kernel.py`
Content:
```
1 #!/usr/bin/env python3
2 import math
3 
4 import torch
5 
6 from ... import settings
7 from ...lazy import KeOpsLazyTensor
8 from .keops_kernel import KeOpsKernel
9 
10 try:
11     from pykeops.torch import LazyTensor as KEOLazyTensor
12 
13     class MaternKernel(KeOpsKernel):
14         """
15         Implements the Matern kernel using KeOps as a driver for kernel matrix multiplies.
16 
17         This class can be used as a drop in replacement for gpytorch.kernels.MaternKernel in most cases, and supports
18         the same arguments. There are currently a few limitations, for example a lack of batch mode support. However,
19         most other features like ARD will work.
20         """
21 
22         has_lengthscale = True
23 
24         def __init__(self, nu=2.5, **kwargs):
25             if nu not in {0.5, 1.5, 2.5}:
26                 raise RuntimeError("nu expected to be 0.5, 1.5, or 2.5")
27             super(MaternKernel, self).__init__(**kwargs)
28             self.nu = nu
29 
30         def _nonkeops_covar_func(self, x1, x2, diag=False):
31             distance = self.covar_dist(x1, x2, diag=diag)
32             exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)
33 
34             if self.nu == 0.5:
35                 constant_component = 1
36             elif self.nu == 1.5:
37                 constant_component = (math.sqrt(3) * distance).add(1)
38             elif self.nu == 2.5:
39                 constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance ** 2)
40             return constant_component * exp_component
41 
42         def covar_func(self, x1, x2, diag=False):
43             # We only should use KeOps on big kernel matrices
44             # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)
45             # then don't apply KeOps
46             if (
47                 diag
48                 or x1.size(-2) < settings.max_cholesky_size.value()
49                 or x2.size(-2) < settings.max_cholesky_size.value()
50             ):
51                 return self._nonkeops_covar_func(x1, x2, diag=diag)
52             # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.
53             # This bug is fixed on KeOps master, and we'll remove that part of the check
54             # when they cut a new release.
55             elif x1.size(-2) == 1 or x2.size(-2) == 1:
56                 return self._nonkeops_covar_func(x1, x2, diag=diag)
57             else:
58                 with torch.autograd.enable_grad():
59                     # We only should use KeOps on big kernel matrices
60                     # If we would otherwise be performing Cholesky inference, then don't apply KeOps
61                     if (
62                         x1.size(-2) < settings.max_cholesky_size.value()
63                         or x2.size(-2) < settings.max_cholesky_size.value()
64                     ):
65                         x1_ = x1[..., :, None, :]
66                         x2_ = x2[..., None, :, :]
67                     else:
68                         x1_ = KEOLazyTensor(x1[..., :, None, :])
69                         x2_ = KEOLazyTensor(x2[..., None, :, :])
70 
71                     distance = ((x1_ - x2_) ** 2).sum(-1).sqrt()
72                     exp_component = (-math.sqrt(self.nu * 2) * distance).exp()
73 
74                     if self.nu == 0.5:
75                         constant_component = 1
76                     elif self.nu == 1.5:
77                         constant_component = (math.sqrt(3) * distance) + 1
78                     elif self.nu == 2.5:
79                         constant_component = (math.sqrt(5) * distance) + (1 + 5.0 / 3.0 * distance ** 2)
80 
81                     return constant_component * exp_component
82 
83         def forward(self, x1, x2, diag=False, **params):
84             mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
85 
86             x1_ = (x1 - mean).div(self.lengthscale)
87             x2_ = (x2 - mean).div(self.lengthscale)
88 
89             if diag:
90                 return self.covar_func(x1_, x2_, diag=True)
91 
92             covar_func = lambda x1, x2, diag=False: self.covar_func(x1, x2, diag)
93             return KeOpsLazyTensor(x1_, x2_, covar_func)
94 
95 
96 except ImportError:
97 
98     class MaternKernel(KeOpsKernel):
99         def __init__(self, *args, **kwargs):
100             super().__init__()
101
```
--- 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/keops/matern_kernel.py b/gpytorch/kernels/keops/matern_kernel.py
--- a/gpytorch/kernels/keops/matern_kernel.py
+++ b/gpytorch/kernels/keops/matern_kernel.py
@@ -43,19 +43,21 @@
             # We only should use KeOps on big kernel matrices
             # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)
             # then don't apply KeOps
-            if (
-                diag
-                or x1.size(-2) < settings.max_cholesky_size.value()
-                or x2.size(-2) < settings.max_cholesky_size.value()
-            ):
-                return self._nonkeops_covar_func(x1, x2, diag=diag)
-            # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.
-            # This bug is fixed on KeOps master, and we'll remove that part of the check
-            # when they cut a new release.
-            elif x1.size(-2) == 1 or x2.size(-2) == 1:
-                return self._nonkeops_covar_func(x1, x2, diag=diag)
-            else:
-                with torch.autograd.enable_grad():
+            # enable gradients to ensure that test time caches on small predictions are still
+            # backprop-able
+            with torch.autograd.enable_grad():
+                if (
+                    diag
+                    or x1.size(-2) < settings.max_cholesky_size.value()
+                    or x2.size(-2) < settings.max_cholesky_size.value()
+                ):
+                    return self._nonkeops_covar_func(x1, x2, diag=diag)
+                # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.
+                # This bug is fixed on KeOps master, and we'll remove that part of the check
+                # when they cut a new release.
+                elif x1.size(-2) == 1 or x2.size(-2) == 1:
+                    return self._nonkeops_covar_func(x1, x2, diag=diag)
+                else:
                     # We only should use KeOps on big kernel matrices
                     # If we would otherwise be performing Cholesky inference, then don't apply KeOps
                     if (
diff --git a/gpytorch/kernels/keops/rbf_kernel.py b/gpytorch/kernels/keops/rbf_kernel.py
--- a/gpytorch/kernels/keops/rbf_kernel.py
+++ b/gpytorch/kernels/keops/rbf_kernel.py
@@ -22,21 +22,28 @@
 
         def _nonkeops_covar_func(self, x1, x2, diag=False):
             return self.covar_dist(
-                x1, x2, square_dist=True, diag=diag, dist_postprocess_func=postprocess_rbf, postprocess=True
+                x1,
+                x2,
+                square_dist=True,
+                diag=diag,
+                dist_postprocess_func=postprocess_rbf,
+                postprocess=True,
             )
 
         def covar_func(self, x1, x2, diag=False):
             # We only should use KeOps on big kernel matrices
             # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)
             # then don't apply KeOps
-            if (
-                diag
-                or x1.size(-2) < settings.max_cholesky_size.value()
-                or x2.size(-2) < settings.max_cholesky_size.value()
-            ):
-                return self._nonkeops_covar_func(x1, x2, diag=diag)
-
+            # enable gradients to ensure that test time caches on small predictions are still
+            # backprop-able
             with torch.autograd.enable_grad():
+                if (
+                    diag
+                    or x1.size(-2) < settings.max_cholesky_size.value()
+                    or x2.size(-2) < settings.max_cholesky_size.value()
+                ):
+                    return self._nonkeops_covar_func(x1, x2, diag=diag)
+
                 x1_ = KEOLazyTensor(x1[..., :, None, :])
                 x2_ = KEOLazyTensor(x2[..., None, :, :])

{"golden_diff": "diff --git a/gpytorch/kernels/keops/matern_kernel.py b/gpytorch/kernels/keops/matern_kernel.py\n--- a/gpytorch/kernels/keops/matern_kernel.py\n+++ b/gpytorch/kernels/keops/matern_kernel.py\n@@ -43,19 +43,21 @@\n             # We only should use KeOps on big kernel matrices\n             # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)\n             # then don't apply KeOps\n-            if (\n-                diag\n-                or x1.size(-2) < settings.max_cholesky_size.value()\n-                or x2.size(-2) < settings.max_cholesky_size.value()\n-            ):\n-                return self._nonkeops_covar_func(x1, x2, diag=diag)\n-            # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.\n-            # This bug is fixed on KeOps master, and we'll remove that part of the check\n-            # when they cut a new release.\n-            elif x1.size(-2) == 1 or x2.size(-2) == 1:\n-                return self._nonkeops_covar_func(x1, x2, diag=diag)\n-            else:\n-                with torch.autograd.enable_grad():\n+            # enable gradients to ensure that test time caches on small predictions are still\n+            # backprop-able\n+            with torch.autograd.enable_grad():\n+                if (\n+                    diag\n+                    or x1.size(-2) < settings.max_cholesky_size.value()\n+                    or x2.size(-2) < settings.max_cholesky_size.value()\n+                ):\n+                    return self._nonkeops_covar_func(x1, x2, diag=diag)\n+                # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.\n+                # This bug is fixed on KeOps master, and we'll remove that part of the check\n+                # when they cut a new release.\n+                elif x1.size(-2) == 1 or x2.size(-2) == 1:\n+                    return self._nonkeops_covar_func(x1, x2, diag=diag)\n+                else:\n                     # We only should use KeOps on big kernel matrices\n                     # If we would otherwise be performing Cholesky inference, then don't apply KeOps\n                     if (\ndiff --git a/gpytorch/kernels/keops/rbf_kernel.py b/gpytorch/kernels/keops/rbf_kernel.py\n--- a/gpytorch/kernels/keops/rbf_kernel.py\n+++ b/gpytorch/kernels/keops/rbf_kernel.py\n@@ -22,21 +22,28 @@\n \n         def _nonkeops_covar_func(self, x1, x2, diag=False):\n             return self.covar_dist(\n-                x1, x2, square_dist=True, diag=diag, dist_postprocess_func=postprocess_rbf, postprocess=True\n+                x1,\n+                x2,\n+                square_dist=True,\n+                diag=diag,\n+                dist_postprocess_func=postprocess_rbf,\n+                postprocess=True,\n             )\n \n         def covar_func(self, x1, x2, diag=False):\n             # We only should use KeOps on big kernel matrices\n             # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)\n             # then don't apply KeOps\n-            if (\n-                diag\n-                or x1.size(-2) < settings.max_cholesky_size.value()\n-                or x2.size(-2) < settings.max_cholesky_size.value()\n-            ):\n-                return self._nonkeops_covar_func(x1, x2, diag=diag)\n-\n+            # enable gradients to ensure that test time caches on small predictions are still\n+            # backprop-able\n             with torch.autograd.enable_grad():\n+                if (\n+                    diag\n+                    or x1.size(-2) < settings.max_cholesky_size.value()\n+                    or x2.size(-2) < settings.max_cholesky_size.value()\n+                ):\n+                    return self._nonkeops_covar_func(x1, x2, diag=diag)\n+\n                 x1_ = KEOLazyTensor(x1[..., :, None, :])\n                 x2_ = KEOLazyTensor(x2[..., None, :, :])\n", "issue": "[Bug] Gradient computation with KeOps kernel only works with low number of training data, fails otherwise\n# \ud83d\udc1b Bug\r\n\r\nHi,\r\n\r\nFirst of all, I would like to extend my thanks to all the developers for all the efforts you have put into both the research and this great package. \r\n\r\nConsider a GP with a KeOps kernel (e.g. gpytorch.kernels.keops.RBFKernel). If I train it with N=100 number of points, then the gradient of predictive mean can be obtained by torch.autograd.grad or .backward(), but set N=500 and an error will be thrown that the input tensor was not used in the graph. I have tested the script on two separate machines and a colab instance. Using the GPyTorch standard kernels will not run into this issue. I spent a good deal of time pinpointing what was wrong from bigger chunks of code and this seemed to be the issue. The gradient link seems to cut off between the covariance output of the KeOps kernel and the input (covar.x1).  \r\n\r\nI have provided a minimal code right below that should quickly give you an idea of this somewhat strange behavior. It contains two test cases with N=100 (which passes) and 500 (fails). The code is from GPyTorch regression examples, I have only added the GP kernel and a few lines for calculating gradients at the end. \r\n\r\n## To reproduce\r\n\r\n```python\r\nimport math\r\nimport torch\r\nimport gpytorch\r\nimport time\r\n\r\n# We will use the simplest form of GP model, exact inference with gpytorch.kernels.keops.RBFKernel\r\nclass ExactGPModel(gpytorch.models.ExactGP):\r\n    def __init__(self, train_x, train_y, likelihood):\r\n        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)\r\n        self.mean_module = gpytorch.means.ConstantMean()\r\n        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.keops.RBFKernel())\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\ndef train_and_eval_GP(N = 100):\r\n    \"\"\"\r\n    inputs:\r\n    N (int): Number of training points\r\n    \"\"\"\r\n    # make train/val/test\r\n    # Training data is 100 points in [0,1] inclusive regularly spaced\r\n    train_x = torch.linspace(0, 1, N)\r\n    # True function is sin(2*pi*x) with Gaussian noise\r\n    train_y = torch.sin(train_x * (2 * math.pi)) + torch.randn(train_x.size()) * math.sqrt(0.04)\r\n    # normalize features\r\n    mean = train_x.mean()\r\n    std = train_x.std() + 1e-6 # prevent dividing by 0\r\n    train_x = (train_x - mean) / std\r\n\r\n    # normalize labels\r\n    mean, std = train_y.mean(),train_y.std()\r\n    train_y = (train_y - mean) / std\r\n\r\n    # make continguous\r\n    train_x, train_y = train_x.contiguous(), train_y.contiguous()\r\n\r\n    output_device = torch.device('cuda:0')\r\n\r\n    train_x, train_y = train_x.to(output_device), train_y.to(output_device)\r\n\r\n    # initialize likelihood and model\r\n    likelihood = gpytorch.likelihoods.GaussianLikelihood().to(output_device)\r\n    model = ExactGPModel(train_x, train_y, likelihood).to(output_device)\r\n\r\n    # Find optimal model hyperparameters\r\n    model.train()\r\n    likelihood.train()\r\n\r\n    # Use the adam optimizer\r\n    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)  # Includes GaussianLikelihood parameters\r\n\r\n    # \"Loss\" for GPs - the marginal log likelihood\r\n    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)\r\n\r\n    training_iter = 20\r\n    for i in range(training_iter):\r\n        # Zero gradients from previous iteration\r\n        optimizer.zero_grad()\r\n        # Output from model\r\n        output = model(train_x)\r\n        # Calc loss and backprop gradients\r\n        loss = -mll(output, train_y)\r\n        loss.backward()\r\n        optimizer.step()\r\n    print(f'GP model trained.')\r\n\r\n    # Get into evaluation (predictive posterior) mode\r\n    model.eval()\r\n    likelihood.eval()\r\n\r\n    # Test points are regularly spaced along [0,1]\r\n    test_x = torch.linspace(0, 1, 51, requires_grad=True).to(output_device).contiguous()\r\n    \r\n    # Make predictions by feeding model through likelihood\r\n    with gpytorch.settings.fast_pred_var():\r\n        observed_pred = likelihood(model(test_x))\r\n        assert torch.autograd.grad(observed_pred.mean.sum(), test_x, retain_graph=True) is not None\r\n        print('gradient exists:')\r\n        print(torch.autograd.grad(observed_pred.mean.sum(), test_x, retain_graph=True))\r\n\r\nif __name__ == \"__main__\":\r\n    Ns = [100, 500] #test cases\r\n    for n in Ns:\r\n        try:\r\n            print(f'testing with {n} points...')\r\n            train_and_eval_GP(N = n) \r\n            print('success!')\r\n        except Exception as e:\r\n            print('failed.')\r\n            print(e)\r\n```\r\n\r\n** Stack trace/error message **\r\n```\r\ntesting with 100 points...\r\nGP model trained.\r\ngradient exists:\r\n(tensor([-2.6629e+00, -2.6507e+00, -2.6344e+00, -2.6138e+00, -2.5891e+00,\r\n        ...,\r\n         9.1844e-01], device='cuda:0'),)\r\nsuccess!\r\ntesting with 500 points...\r\nGP model trained.\r\nfailed.\r\nOne of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.\r\n```\r\n\r\n## Expected Behavior\r\n\r\nSetting the kernel to the standard, non-KeOps kernel (gpytorch.kernels.RBFKernel) we get the gradients for the second case, shown below. However, I can't simply use it since I'm working on a larger dataset that will run out of memory if I do so.\r\n\r\n```\r\ntesting with 100 points...\r\nGP model trained.\r\ngradient exists:\r\n(tensor([-2.5885, -2.5870, -2.5819, -2.5732, -2.5609, -2.5449, -2.5254, -2.5022,\r\n        ...,\r\n         0.6977,  0.7880,  0.8762], device='cuda:0'),)\r\nsuccess!\r\ntesting with 500 points...\r\nGP model trained.\r\ngradient exists:\r\n(tensor([-2.4751, -2.4741, -2.4698, -2.4622, -2.4514, -2.4373, -2.4199, -2.3993,\r\n        ...,\r\n         0.4727,  0.5600,  0.6462], device='cuda:0'),)\r\nsuccess!\r\n```\r\n## System information\r\n\r\n**Please complete the following information:**\r\n- <!-- GPyTorch Version (run `print(gpytorch.__version__)` --> GPyTorch 1.6.0\r\n- <!-- PyTorch Version (run `print(torch.__version__)` --> PyTorch  1.10.1\r\n- <!-- KeOps Version --> KeOps 1.5\r\n- <!-- Computer OS --> Ubuntu 20.04.3 LTS\r\n\r\n## Additional context\r\nI know that the issue might be unrelated to GPyTorch as this clearly stems from the keops kernel. It is however difficult to track, so I thought I'd report it here. \r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\nimport torch\n\nfrom ... import settings\nfrom ...lazy import KeOpsLazyTensor\nfrom ..rbf_kernel import postprocess_rbf\nfrom .keops_kernel import KeOpsKernel\n\ntry:\n    from pykeops.torch import LazyTensor as KEOLazyTensor\n\n    class RBFKernel(KeOpsKernel):\n        \"\"\"\n        Implements the RBF kernel using KeOps as a driver for kernel matrix multiplies.\n\n        This class can be used as a drop in replacement for gpytorch.kernels.RBFKernel in most cases, and supports\n        the same arguments. There are currently a few limitations, for example a lack of batch mode support. However,\n        most other features like ARD will work.\n        \"\"\"\n\n        has_lengthscale = True\n\n        def _nonkeops_covar_func(self, x1, x2, diag=False):\n            return self.covar_dist(\n                x1, x2, square_dist=True, diag=diag, dist_postprocess_func=postprocess_rbf, postprocess=True\n            )\n\n        def covar_func(self, x1, x2, diag=False):\n            # We only should use KeOps on big kernel matrices\n            # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)\n            # then don't apply KeOps\n            if (\n                diag\n                or x1.size(-2) < settings.max_cholesky_size.value()\n                or x2.size(-2) < settings.max_cholesky_size.value()\n            ):\n                return self._nonkeops_covar_func(x1, x2, diag=diag)\n\n            with torch.autograd.enable_grad():\n                x1_ = KEOLazyTensor(x1[..., :, None, :])\n                x2_ = KEOLazyTensor(x2[..., None, :, :])\n\n                K = (-((x1_ - x2_) ** 2).sum(-1) / 2).exp()\n\n                return K\n\n        def forward(self, x1, x2, diag=False, **params):\n            x1_ = x1.div(self.lengthscale)\n            x2_ = x2.div(self.lengthscale)\n\n            covar_func = lambda x1, x2, diag=diag: self.covar_func(x1, x2, diag)\n\n            if diag:\n                return covar_func(x1_, x2_, diag=True)\n\n            return KeOpsLazyTensor(x1_, x2_, covar_func)\n\n\nexcept ImportError:\n\n    class RBFKernel(KeOpsKernel):\n        def __init__(self, *args, **kwargs):\n            super().__init__()\n", "path": "gpytorch/kernels/keops/rbf_kernel.py"}, {"content": "#!/usr/bin/env python3\nimport math\n\nimport torch\n\nfrom ... import settings\nfrom ...lazy import KeOpsLazyTensor\nfrom .keops_kernel import KeOpsKernel\n\ntry:\n    from pykeops.torch import LazyTensor as KEOLazyTensor\n\n    class MaternKernel(KeOpsKernel):\n        \"\"\"\n        Implements the Matern kernel using KeOps as a driver for kernel matrix multiplies.\n\n        This class can be used as a drop in replacement for gpytorch.kernels.MaternKernel in most cases, and supports\n        the same arguments. There are currently a few limitations, for example a lack of batch mode support. However,\n        most other features like ARD will work.\n        \"\"\"\n\n        has_lengthscale = True\n\n        def __init__(self, nu=2.5, **kwargs):\n            if nu not in {0.5, 1.5, 2.5}:\n                raise RuntimeError(\"nu expected to be 0.5, 1.5, or 2.5\")\n            super(MaternKernel, self).__init__(**kwargs)\n            self.nu = nu\n\n        def _nonkeops_covar_func(self, x1, x2, diag=False):\n            distance = self.covar_dist(x1, x2, diag=diag)\n            exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)\n\n            if self.nu == 0.5:\n                constant_component = 1\n            elif self.nu == 1.5:\n                constant_component = (math.sqrt(3) * distance).add(1)\n            elif self.nu == 2.5:\n                constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance ** 2)\n            return constant_component * exp_component\n\n        def covar_func(self, x1, x2, diag=False):\n            # We only should use KeOps on big kernel matrices\n            # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)\n            # then don't apply KeOps\n            if (\n                diag\n                or x1.size(-2) < settings.max_cholesky_size.value()\n                or x2.size(-2) < settings.max_cholesky_size.value()\n            ):\n                return self._nonkeops_covar_func(x1, x2, diag=diag)\n            # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.\n            # This bug is fixed on KeOps master, and we'll remove that part of the check\n            # when they cut a new release.\n            elif x1.size(-2) == 1 or x2.size(-2) == 1:\n                return self._nonkeops_covar_func(x1, x2, diag=diag)\n            else:\n                with torch.autograd.enable_grad():\n                    # We only should use KeOps on big kernel matrices\n                    # If we would otherwise be performing Cholesky inference, then don't apply KeOps\n                    if (\n                        x1.size(-2) < settings.max_cholesky_size.value()\n                        or x2.size(-2) < settings.max_cholesky_size.value()\n                    ):\n                        x1_ = x1[..., :, None, :]\n                        x2_ = x2[..., None, :, :]\n                    else:\n                        x1_ = KEOLazyTensor(x1[..., :, None, :])\n                        x2_ = KEOLazyTensor(x2[..., None, :, :])\n\n                    distance = ((x1_ - x2_) ** 2).sum(-1).sqrt()\n                    exp_component = (-math.sqrt(self.nu * 2) * distance).exp()\n\n                    if self.nu == 0.5:\n                        constant_component = 1\n                    elif self.nu == 1.5:\n                        constant_component = (math.sqrt(3) * distance) + 1\n                    elif self.nu == 2.5:\n                        constant_component = (math.sqrt(5) * distance) + (1 + 5.0 / 3.0 * distance ** 2)\n\n                    return constant_component * exp_component\n\n        def forward(self, x1, x2, diag=False, **params):\n            mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]\n\n            x1_ = (x1 - mean).div(self.lengthscale)\n            x2_ = (x2 - mean).div(self.lengthscale)\n\n            if diag:\n                return self.covar_func(x1_, x2_, diag=True)\n\n            covar_func = lambda x1, x2, diag=False: self.covar_func(x1, x2, diag)\n            return KeOpsLazyTensor(x1_, x2_, covar_func)\n\n\nexcept ImportError:\n\n    class MaternKernel(KeOpsKernel):\n        def __init__(self, *args, **kwargs):\n            super().__init__()\n", "path": "gpytorch/kernels/keops/matern_kernel.py"}], "after_files": [{"content": "#!/usr/bin/env python3\nimport torch\n\nfrom ... import settings\nfrom ...lazy import KeOpsLazyTensor\nfrom ..rbf_kernel import postprocess_rbf\nfrom .keops_kernel import KeOpsKernel\n\ntry:\n    from pykeops.torch import LazyTensor as KEOLazyTensor\n\n    class RBFKernel(KeOpsKernel):\n        \"\"\"\n        Implements the RBF kernel using KeOps as a driver for kernel matrix multiplies.\n\n        This class can be used as a drop in replacement for gpytorch.kernels.RBFKernel in most cases, and supports\n        the same arguments. There are currently a few limitations, for example a lack of batch mode support. However,\n        most other features like ARD will work.\n        \"\"\"\n\n        has_lengthscale = True\n\n        def _nonkeops_covar_func(self, x1, x2, diag=False):\n            return self.covar_dist(\n                x1,\n                x2,\n                square_dist=True,\n                diag=diag,\n                dist_postprocess_func=postprocess_rbf,\n                postprocess=True,\n            )\n\n        def covar_func(self, x1, x2, diag=False):\n            # We only should use KeOps on big kernel matrices\n            # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)\n            # then don't apply KeOps\n            # enable gradients to ensure that test time caches on small predictions are still\n            # backprop-able\n            with torch.autograd.enable_grad():\n                if (\n                    diag\n                    or x1.size(-2) < settings.max_cholesky_size.value()\n                    or x2.size(-2) < settings.max_cholesky_size.value()\n                ):\n                    return self._nonkeops_covar_func(x1, x2, diag=diag)\n\n                x1_ = KEOLazyTensor(x1[..., :, None, :])\n                x2_ = KEOLazyTensor(x2[..., None, :, :])\n\n                K = (-((x1_ - x2_) ** 2).sum(-1) / 2).exp()\n\n                return K\n\n        def forward(self, x1, x2, diag=False, **params):\n            x1_ = x1.div(self.lengthscale)\n            x2_ = x2.div(self.lengthscale)\n\n            covar_func = lambda x1, x2, diag=diag: self.covar_func(x1, x2, diag)\n\n            if diag:\n                return covar_func(x1_, x2_, diag=True)\n\n            return KeOpsLazyTensor(x1_, x2_, covar_func)\n\n\nexcept ImportError:\n\n    class RBFKernel(KeOpsKernel):\n        def __init__(self, *args, **kwargs):\n            super().__init__()\n", "path": "gpytorch/kernels/keops/rbf_kernel.py"}, {"content": "#!/usr/bin/env python3\nimport math\n\nimport torch\n\nfrom ... import settings\nfrom ...lazy import KeOpsLazyTensor\nfrom .keops_kernel import KeOpsKernel\n\ntry:\n    from pykeops.torch import LazyTensor as KEOLazyTensor\n\n    class MaternKernel(KeOpsKernel):\n        \"\"\"\n        Implements the Matern kernel using KeOps as a driver for kernel matrix multiplies.\n\n        This class can be used as a drop in replacement for gpytorch.kernels.MaternKernel in most cases, and supports\n        the same arguments. There are currently a few limitations, for example a lack of batch mode support. However,\n        most other features like ARD will work.\n        \"\"\"\n\n        has_lengthscale = True\n\n        def __init__(self, nu=2.5, **kwargs):\n            if nu not in {0.5, 1.5, 2.5}:\n                raise RuntimeError(\"nu expected to be 0.5, 1.5, or 2.5\")\n            super(MaternKernel, self).__init__(**kwargs)\n            self.nu = nu\n\n        def _nonkeops_covar_func(self, x1, x2, diag=False):\n            distance = self.covar_dist(x1, x2, diag=diag)\n            exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)\n\n            if self.nu == 0.5:\n                constant_component = 1\n            elif self.nu == 1.5:\n                constant_component = (math.sqrt(3) * distance).add(1)\n            elif self.nu == 2.5:\n                constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance ** 2)\n            return constant_component * exp_component\n\n        def covar_func(self, x1, x2, diag=False):\n            # We only should use KeOps on big kernel matrices\n            # If we would otherwise be performing Cholesky inference, (or when just computing a kernel matrix diag)\n            # then don't apply KeOps\n            # enable gradients to ensure that test time caches on small predictions are still\n            # backprop-able\n            with torch.autograd.enable_grad():\n                if (\n                    diag\n                    or x1.size(-2) < settings.max_cholesky_size.value()\n                    or x2.size(-2) < settings.max_cholesky_size.value()\n                ):\n                    return self._nonkeops_covar_func(x1, x2, diag=diag)\n                # TODO: x1 / x2 size checks are a work around for a very minor bug in KeOps.\n                # This bug is fixed on KeOps master, and we'll remove that part of the check\n                # when they cut a new release.\n                elif x1.size(-2) == 1 or x2.size(-2) == 1:\n                    return self._nonkeops_covar_func(x1, x2, diag=diag)\n                else:\n                    # We only should use KeOps on big kernel matrices\n                    # If we would otherwise be performing Cholesky inference, then don't apply KeOps\n                    if (\n                        x1.size(-2) < settings.max_cholesky_size.value()\n                        or x2.size(-2) < settings.max_cholesky_size.value()\n                    ):\n                        x1_ = x1[..., :, None, :]\n                        x2_ = x2[..., None, :, :]\n                    else:\n                        x1_ = KEOLazyTensor(x1[..., :, None, :])\n                        x2_ = KEOLazyTensor(x2[..., None, :, :])\n\n                    distance = ((x1_ - x2_) ** 2).sum(-1).sqrt()\n                    exp_component = (-math.sqrt(self.nu * 2) * distance).exp()\n\n                    if self.nu == 0.5:\n                        constant_component = 1\n                    elif self.nu == 1.5:\n                        constant_component = (math.sqrt(3) * distance) + 1\n                    elif self.nu == 2.5:\n                        constant_component = (math.sqrt(5) * distance) + (1 + 5.0 / 3.0 * distance ** 2)\n\n                    return constant_component * exp_component\n\n        def forward(self, x1, x2, diag=False, **params):\n            mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]\n\n            x1_ = (x1 - mean).div(self.lengthscale)\n            x2_ = (x2 - mean).div(self.lengthscale)\n\n            if diag:\n                return self.covar_func(x1_, x2_, diag=True)\n\n            covar_func = lambda x1, x2, diag=False: self.covar_func(x1, x2, diag)\n            return KeOpsLazyTensor(x1_, x2_, covar_func)\n\n\nexcept ImportError:\n\n    class MaternKernel(KeOpsKernel):\n        def __init__(self, *args, **kwargs):\n            super().__init__()\n", "path": "gpytorch/kernels/keops/matern_kernel.py"}]}