Repository: cornellius-gp/gpytorch
Issue: Error when using `fast_pred_var` with only 1 training input

I'm getting an error when I try to use `fast_pred_var` but not when I do predictions without this setting.

To reproduce:

```python
import torch
import gpytorch

# from the gp regression tutorial
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.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)

n_train = 1
n_test = 100
n_features = 3
    
likelihood = gpytorch.likelihoods.GaussianLikelihood()
train_x = torch.rand(n_train, n_features)
train_y = torch.rand(n_train)
model = ExactGPModel(train_x, train_y, likelihood)

model.eval()
model.likelihood.eval()

# this works fine
model(torch.rand(n_test, n_features))

# this throws the error
with gpytorch.settings.fast_pred_var():
    model(torch.rand(n_test, n_features))
```

The error that I get is

```python
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-1-05bfb3c2646a> in <module>
     27 # this throws the error
     28 with gpytorch.settings.fast_pred_var():
---> 29     model(torch.rand(100, 2))

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
    262 
    263             predictive_mean = self.prediction_strategy.exact_predictive_mean(test_mean, test_train_covar)
--> 264             predictive_covar = self.prediction_strategy.exact_predictive_covar(test_test_covar, test_train_covar)
    265 
    266             if num_tasks > 1:

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/models/exact_prediction_strategies.py in exact_predictive_covar(self, test_test_covar, test_train_covar)
    324             return test_test_covar + MatmulLazyTensor(test_train_covar, covar_correction_rhs)
    325 
--> 326         precomputed_cache = self.covar_cache
    327         covar_inv_quad_form_root = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache,
    328                                                                                    test_train_covar)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
     15         cache_name = name if name is not None else method
     16         if cache_name not in self._memoize_cache:
---> 17             self._memoize_cache[cache_name] = method(self, *args, **kwargs)
     18         return self._memoize_cache[cache_name]
     19 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/models/exact_prediction_strategies.py in covar_cache(self)
    265             train_train_covar_inv_root = train_train_covar[0].root_inv_decomposition().root.evaluate()
    266         else:
--> 267             train_train_covar_inv_root = train_train_covar.root_inv_decomposition().root.evaluate()
    268 
    269         return self._exact_predictive_covar_inv_quad_form_cache(train_train_covar_inv_root, self._last_test_train_covar)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
     15         cache_name = name if name is not None else method
     16         if cache_name not in self._memoize_cache:
---> 17             self._memoize_cache[cache_name] = method(self, *args, **kwargs)
     18         return self._memoize_cache[cache_name]
     19 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in root_inv_decomposition(self, initial_vectors, test_vectors)
   1088             inverse=True,
   1089             initial_vectors=initial_vectors,
-> 1090         )(*self.representation())
   1091 
   1092         if initial_vectors is not None and initial_vectors.size(-1) > 1:

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/functions/_root_decomposition.py in forward(self, *matrix_args)
     51             matrix_shape=self.matrix_shape,
     52             batch_shape=self.batch_shape,
---> 53             init_vecs=self.initial_vectors,
     54         )
     55 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/lib/python3.7/site-packages/gpytorch/utils/lanczos.py in lanczos_tridiag(matmul_closure, max_iter, dtype, device, matrix_shape, batch_shape, init_vecs, num_init_vecs, tol)
     80     # Copy over alpha_0 and beta_0 to t_mat
     81     t_mat[0, 0].copy_(alpha_0)
---> 82     t_mat[0, 1].copy_(beta_0)
     83     t_mat[1, 0].copy_(beta_0)
     84 

IndexError: index 1 is out of bounds for dimension 0 with size 1
```

I played with this a little bit, and it's only come up for me when I have just a single training input. For now, I just have a workaround where I make sure not to use `fast_pred_var` unless there are multiple training inputs, but it would be great if this always worked.

GPytorch version: 0.2.1
OS: Ubuntu 16.04
Python: 3.7

Changed files:
gpytorch/lazy/lazy_tensor.py, test/examples/test_simple_gp_regression.py

diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 82d6d15b..5e2b6adc 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -1065,6 +1065,9 @@ class LazyTensor(object):
         """
         from .root_lazy_tensor import RootLazyTensor
 
+        if self.shape[-2:].numel() == 1:
+            return RootLazyTensor(1 / self.evaluate().sqrt())
+
         if not self.is_square:
             raise RuntimeError(
                 "root_inv_decomposition only operates on (batches of) square (symmetric) LazyTensors. "
diff --git a/test/examples/test_simple_gp_regression.py b/test/examples/test_simple_gp_regression.py
index e68a06e0..892ecd99 100644
--- a/test/examples/test_simple_gp_regression.py
+++ b/test/examples/test_simple_gp_regression.py
@@ -185,6 +185,34 @@ class TestSimpleGPRegression(unittest.TestCase):
             self.assertTrue(torch.allclose(mean_skip_var, mean))
             self.assertTrue(torch.allclose(mean_skip_var, likelihood_mean))
 
+    def test_gp_posterior_single_training_point_smoke_test(self):
+        train_x, test_x, train_y, _ = self._get_data()
+        train_x = train_x[0].unsqueeze(-1).unsqueeze(-1)
+        train_y = train_y[0].unsqueeze(-1)
+        likelihood = GaussianLikelihood()
+        gp_model = ExactGPModel(train_x, train_y, likelihood)
+
+        gp_model.eval()
+        likelihood.eval()
+
+        with gpytorch.settings.fast_pred_var():
+            preds = gp_model(test_x)
+            single_mean = preds.mean
+            single_variance = preds.variance
+
+        self.assertFalse(torch.any(torch.isnan(single_variance)))
+        self.assertFalse(torch.any(torch.isnan(single_mean)))
+
+        gp_model.train()
+        gp_model.eval()
+
+        preds = gp_model(test_x)
+        single_mean = preds.mean
+        single_variance = preds.variance
+
+        self.assertFalse(torch.any(torch.isnan(single_variance)))
+        self.assertFalse(torch.any(torch.isnan(single_mean)))
+
     def test_posterior_latent_gp_and_likelihood_with_optimization(self, cuda=False, checkpoint=0):
         train_x, test_x, train_y, test_y = self._get_data(
             cuda=cuda, num_data=(1000 if checkpoint else 11), add_noise=bool(checkpoint),