Repository: cornellius-gp/gpytorch
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. 


Changed files:
gpytorch/kernels/keops/rbf_kernel.py, gpytorch/kernels/keops/matern_kernel.py

diff --git a/gpytorch/kernels/keops/matern_kernel.py b/gpytorch/kernels/keops/matern_kernel.py
index 4abd7f39..84a2b7d9 100644
--- a/gpytorch/kernels/keops/matern_kernel.py
+++ b/gpytorch/kernels/keops/matern_kernel.py
@@ -43,19 +43,21 @@ try:
             # 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
index 163986ee..aa2510ec 100644
--- a/gpytorch/kernels/keops/rbf_kernel.py
+++ b/gpytorch/kernels/keops/rbf_kernel.py
@@ -22,21 +22,28 @@ try:
 
         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, :, :])