Repository: cornellius-gp/gpytorch
Issue: [Bug]

# 🐛 Bug

Batched fixed noise GPs fail when the preconditioning threshold is reached. The concatenation in this line fails with a shape error.

https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/lazy/added_diag_lazy_tensor.py#L126

## To reproduce

** Code snippet to reproduce **
```python
import torch
import gpytorch

class BatchFixedNoiseGP(gpytorch.models.GP):
    def __init__(self, init_x, init_y, noise, batch_shape):
        super().__init__()
        self.mean_module = gpytorch.means.ZeroMean()
        self.covar_module = gpytorch.kernels.RBFKernel(batch_shape=batch_shape)
        self.likelihood = gpytorch.likelihoods.FixedNoiseGaussianLikelihood(noise)
        
    def forward(self, inputs):
        mean = self.mean_module(inputs)
        covar = self.covar_module(inputs)
        return gpytorch.distributions.MultivariateNormal(mean, covar)

batch_shape = [2]
train_x = torch.randn(*batch_shape, 101, 3)
train_y = torch.randn(*batch_shape, 101)
train_noise = torch.rand(*batch_shape, 101)

gp = BatchFixedNoiseGP(train_x, train_y, train_noise, batch_shape)
mll = gpytorch.mlls.ExactMarginalLogLikelihood(gp.likelihood, gp)

with gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):
    train_dist = gp(train_x)
    loss = -mll(train_dist, train_y).sum()
```

** Stack trace/error message **
```
RuntimeError                              Traceback (most recent call last)
<ipython-input-4-9e151e2de37a> in <module>
     24 with gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):
     25     train_dist = gp(train_x)
---> 26     loss = -mll(train_dist, train_y).sum()

~/Code/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     26 
     27     def __call__(self, *inputs, **kwargs):
---> 28         outputs = self.forward(*inputs, **kwargs)
     29         if isinstance(outputs, list):
     30             return [_validate_module_outputs(output) for output in outputs]

~/Code/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, function_dist, target, *params)
     49         # Get the log prob of the marginal distribution
     50         output = self.likelihood(function_dist, *params)
---> 51         res = output.log_prob(target)
     52 
     53         # Add additional terms (SGPR / learned inducing points, heteroskedastic likelihood models)

~/Code/gpytorch/gpytorch/distributions/multivariate_normal.py in log_prob(self, value)
    138 
    139         # Get log determininat and first part of quadratic form
--> 140         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    141 
    142         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])

~/Code/gpytorch/gpytorch/lazy/lazy_tensor.py in inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
   1069             probe_vectors,
   1070             probe_vector_norms,
-> 1071             *args,
   1072         )
   1073 

~/Code/gpytorch/gpytorch/functions/_inv_quad_log_det.py in forward(ctx, representation_tree, dtype, device, matrix_shape, batch_shape, inv_quad, logdet, probe_vectors, probe_vector_norms, *args)
     65         lazy_tsr = ctx.representation_tree(*matrix_args)
     66         with torch.no_grad():
---> 67             preconditioner, precond_lt, logdet_correction = lazy_tsr._preconditioner()
     68 
     69         ctx.preconditioner = preconditioner

~/Code/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _preconditioner(self)
     84                 )
     85                 return None, None, None
---> 86             self._init_cache()
     87 
     88         # NOTE: We cannot memoize this precondition closure as it causes a memory leak

~/Code/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache(self)
    107             self._init_cache_for_constant_diag(eye, batch_shape, n, k)
    108         else:
--> 109             self._init_cache_for_non_constant_diag(eye, batch_shape, n)
    110 
    111         self._precond_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)

~/Code/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache_for_non_constant_diag(self, eye, batch_shape, n)
    125         # With non-constant diagonals, we cant factor out the noise as easily
    126         # eye = eye.expand(*batch_shape, -1, -1)
--> 127         self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))
    128         self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()
    129 

RuntimeError: Tensors must have same number of dimensions: got 3 and 2
```

## Expected Behavior

Everything works fine until the preconditioning threshold is reached. Obviously one would hope that it would continue to work.

## System information

**Please complete the following information:**
GPyTorch Version: 1.2.0
PyTorch Version:  '1.6.0.dev20200522'
OS: Ubuntu 16.04 LTS

## Additional context
This appears to fix the problem
```
    def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):
        # With non-constant diagonals, we cant factor out the noise as easily
        eye = eye.expand(*batch_shape, -1, -1)
        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye), dim=-2))
```


Changed files:
gpytorch/lazy/added_diag_lazy_tensor.py, test/examples/test_white_noise_regression.py

diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py
index 38489f5d..729a7955 100644
--- a/gpytorch/lazy/added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/added_diag_lazy_tensor.py
@@ -68,12 +68,33 @@ class AddedDiagLazyTensor(SumLazyTensor):
             return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)
 
     def _preconditioner(self):
+        r"""
+        Here we use a partial pivoted Cholesky preconditioner:
+
+        K \\approx L L^T + D
+
+        where L L^T is a low rank approximation, and D is a diagonal.
+        We can compute the preconditioner's inverse using Woodbury
+
+        (L L^T + D)^{-1} = D^{-1} - D^{-1} L (I + L D^{-1} L^T)^{-1} L^T D^{-1}
+
+        This function returns:
+        - A function `precondition_closure` that computes the solve (L L^T + D)^{-1} x
+        - A LazyTensor `precondition_lt` that represents (L L^T + D)
+        - The log determinant of (L L^T + D)
+        """
+
         if self.preconditioner_override is not None:
             return self.preconditioner_override(self)
 
         if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():
             return None, None, None
 
+        # Cache a QR decomposition [Q; Q'] R = [D^{-1/2}; L]
+        # This makes it fast to compute solves and log determinants with it
+        #
+        # Through woodbury, (L L^T + D)^{-1} reduces down to (D^{-1} - D^{-1/2} Q Q^T D^{-1/2})
+        # Through matrix determinant lemma, log |L L^T + D| reduces down to 2 log |R|
         if self._q_cache is None:
             max_iter = settings.max_preconditioner_size.value()
             self._piv_chol_self = pivoted_cholesky.pivoted_cholesky(self._lazy_tensor, max_iter)
@@ -87,6 +108,7 @@ class AddedDiagLazyTensor(SumLazyTensor):
 
         # NOTE: We cannot memoize this precondition closure as it causes a memory leak
         def precondition_closure(tensor):
+            # This makes it fast to compute solves with it
             qqt = self._q_cache.matmul(self._q_cache.transpose(-2, -1).matmul(tensor))
             if self._constant_diag:
                 return (1 / self._noise) * (tensor - qqt)
@@ -102,6 +124,7 @@ class AddedDiagLazyTensor(SumLazyTensor):
         noise_first_element = self._noise[..., :1, :]
         self._constant_diag = torch.equal(self._noise, noise_first_element * torch.ones_like(self._noise))
         eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)
+        eye = eye.expand(*batch_shape, k, k)
 
         if self._constant_diag:
             self._init_cache_for_constant_diag(eye, batch_shape, n, k)
@@ -123,9 +146,10 @@ class AddedDiagLazyTensor(SumLazyTensor):
 
     def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):
         # With non-constant diagonals, we cant factor out the noise as easily
-        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))
+        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye), dim=-2))
         self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()
 
+        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
         logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
         logdet -= (1.0 / self._noise).log().sum([-1, -2])
         self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
diff --git a/test/examples/test_white_noise_regression.py b/test/examples/test_white_noise_regression.py
index 2b9cc22e..30a438c6 100644
--- a/test/examples/test_white_noise_regression.py
+++ b/test/examples/test_white_noise_regression.py
@@ -53,43 +53,7 @@ class TestWhiteNoiseGPRegression(unittest.TestCase):
         test_y = torch.sin(test_x * (2 * pi))
         return train_x, test_x, train_y, test_y
 
-    def test_posterior_latent_gp_and_likelihood_without_optimization(self, cuda=False):
-        train_x, test_x, train_y, test_y = self._get_data(cuda=cuda)
-        with gpytorch.settings.debug(False):
-            # We're manually going to set the hyperparameters to be ridiculous
-            likelihood = FixedNoiseGaussianLikelihood(torch.ones(11) * 1e-8)
-            gp_model = ExactGPModel(train_x, train_y, likelihood)
-            # Update lengthscale prior to accommodate extreme parameters
-            gp_model.rbf_covar_module.initialize(lengthscale=exp(-6))
-            gp_model.mean_module.initialize(constant=0)
-
-            if cuda:
-                gp_model.cuda()
-                likelihood.cuda()
-
-            # Compute posterior distribution
-            gp_model.eval()
-            likelihood.eval()
-
-            # Let's see how our model does, conditioned with weird hyperparams
-            # The posterior should fit all the data
-            function_predictions = likelihood(gp_model(train_x))
-
-            self.assertLess(torch.norm(function_predictions.mean - train_y), 1e-3)
-            self.assertLess(torch.norm(function_predictions.variance), 5e-3)
-
-            # It shouldn't fit much else though
-            test_function_predictions = gp_model(torch.tensor([1.1]).type_as(test_x))
-
-            self.assertLess(torch.norm(test_function_predictions.mean - 0), 1e-4)
-            self.assertLess(torch.norm(test_function_predictions.variance - gp_model.covar_module.outputscale), 1e-4)
-
-    def test_posterior_latent_gp_and_likelihood_without_optimization_cuda(self):
-        if torch.cuda.is_available():
-            with least_used_cuda_device():
-                self.test_posterior_latent_gp_and_likelihood_without_optimization(cuda=True)
-
-    def test_posterior_latent_gp_and_likelihood_with_optimization(self, cuda=False):
+    def test_posterior_latent_gp_and_likelihood_with_optimization(self, cuda=False, cg=False):
         # This test throws a warning because the fixed noise likelihood gets the wrong input
         warnings.simplefilter("ignore", GPInputWarning)
 
@@ -111,7 +75,7 @@ class TestWhiteNoiseGPRegression(unittest.TestCase):
 
         optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
         optimizer.n_iter = 0
-        with gpytorch.settings.debug(False):
+        with gpytorch.settings.debug(False), gpytorch.settings.max_cholesky_size(0 if cg else 1e10):
             for _ in range(75):
                 optimizer.zero_grad()
                 output = gp_model(train_x)
@@ -133,6 +97,9 @@ class TestWhiteNoiseGPRegression(unittest.TestCase):
 
         self.assertLess(mean_abs_error.squeeze().item(), 0.05)
 
+    def test_posterior_latent_gp_and_likelihood_with_optimization_with_cg(self):
+        return self.test_posterior_latent_gp_and_likelihood_with_optimization(cg=True)
+
     def test_posterior_latent_gp_and_likelihood_with_optimization_cuda(self):
         if torch.cuda.is_available():
             with least_used_cuda_device():