Repository: cornellius-gp/gpytorch

Issue:
KeOps prediction performance
Hi there - really excited about the KeOps integration.  Are there any specific gpytorch settings to help with the prediction speed performance here?  I tried using it as a drop in for other modules I had, and any prediction thats not 1d is taking a significantly longer time.  I realize this is a new feature but the time spent on prediction is a major bottleneck, to that point that it's making me wonder if I have the wrong installs.  One thing I noticed is splitting the batch predictions into chunks of say 1024 speeds up the total compilation and prediction time.  That still makes a variation gp prediction that previously took 30 minutes now a tqdm-estimated 30 hours.  I'm trying to reproduce the KeOps_GP_Regression.ipynb on google colab and it's been stuck on the prediction part for over half an hour.  Install settings are:
`
!pip install pykeops[full]
`
`
!pip install --upgrade --force-reinstall git+https://github.com/cornellius-gp/gpytorch.git@487b90048b83134be50b802e5edb9fefbda6c1cf # specific working commit for stability
`

Leaving this hanging indefinitely:
```
%%time
model.eval()
likelihood.eval()
with torch.no_grad(), gpytorch.settings.fast_pred_var():
    observed_pred = likelihood(model(test_x))
```


Discussion:
@stanbiryukov At the moment, KeOps is significantly slower for variational models (which you mention) than native PyTorch is -- this is caused by a known issue on their end that @jeanfeydy is working on. I assume you aren't also seeing these problems for exact GPs? If so, we'd need more information, because we've tested the integration fairly extensively for exact GPs on fairly large datasets and didn't see any performance issues.

Given the rather extreme speed ups we see on exact GPs, I'm hopeful that we can eventually see some of that speed up in variational models as well. Once these issues appear to be resolved, I will be sure to create an example notebook of the integration with variational / deep GP models.
Good to know regarding variational models.  I'm specifically looking to use exact gp via keops in cases where I previously had to use variational models.  The predictions are much more accurate for the small test cases I've ran.  I had a model running for hours just trying to compile the prediction kernel for ~40k points.  Essentially I'm wondering if there's an optimal minima between compile time and size of tensor to feed into the keops-based kernel for model prediction.  As an example, I'm trying to reproduce the smaller uci-protein example you had and google colab is indefinitely stuck on the full test_x prediction, whereas the prediction on 1024 rows finished in about a minute. 
https://colab.research.google.com/drive/1mckwO1XqGA0NFzOw7qvYDFqpCDoGrXG1

Resolving patch:
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -309,13 +309,11 @@ def exact_predictive_mean(self, test_mean, test_train_covar):
         Returns:
             :obj:`torch.tensor`: The predictive posterior mean of the test points
         """
-        # For efficiency - we can use addmv in the 2d case
-        if test_train_covar.dim() == 2:
-            res = torch.addmv(test_mean, delazify(test_train_covar), self.mean_cache)
-        # In other cases - we'll use the standard infrastructure
-        else:
-            res = (test_train_covar @ self.mean_cache.unsqueeze(-1)).squeeze(-1)
-            res = res + test_mean
+        # NOTE TO FUTURE SELF:
+        # You **cannot* use addmv here, because test_train_covar may not actually be a non lazy tensor even for an exact
+        # GP, and using addmv requires you to delazify test_train_covar, which is obviously a huge no-no!
+        res = (test_train_covar @ self.mean_cache.unsqueeze(-1)).squeeze(-1)
+        res = res + test_mean
         return res
 
     def exact_predictive_covar(self, test_test_covar, test_train_covar):


Test patch: