Repository: cornellius-gp/gpytorch

Issue:
Batched Multi-Output GP preconditioner doesn't compute noise model properly[Bug]
# 🐛 Bug

The noise checking to see if the diagonal is constant in the `init_cache` method for the preconditioner is incorrect for batched GPs. Specifically, [here](https://github.com/cornellius-gp/gpytorch/blob/1e19a641d301218936008f828bfb6e8da081ad3d/gpytorch/lazy/added_diag_lazy_tensor.py#L96) indexes the noise improperly when the noise could be batched.

## To reproduce

** Code snippet to reproduce **
```python
import torch
import gpytorch

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, batch_shape=torch.Size([])):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=batch_shape)
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(batch_shape=batch_shape),
                                                        batch_shape=batch_shape)
    
    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

## base case that works
train_x = torch.randn(301, 1)
train_y = torch.sin(3.1 * train_x) + 0.1 * torch.randn_like(train_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = ExactGPModel(train_x, train_y, likelihood) 

with gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):

    print(likelihood(model(train_x)).log_prob(train_y).sum())

### This doesn't work
### I'm playing with the settings here to force it to crash but it's a shaping error.
train_x = torch.randn(301, 1)
train_y = torch.cat((torch.sin(3.1 * train_x) + 0.1 * torch.randn_like(train_x),
                     torch.sin(train_x) + 0.1 * torch.randn_like(train_x)),dim=1)

likelihood = gpytorch.likelihoods.GaussianLikelihood(batch_shape=torch.Size([2]))
model = ExactGPModel(train_x, train_y, likelihood, batch_shape=torch.Size([2])) 

optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

for i in range(3):
    optimizer.zero_grad()

    with gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):

        loss = -likelihood(model(train_x)).log_prob(train_y.t()).sum()
        loss.backward()
        optimizer.step()
```

** Stack trace/error message **
The noise checking is actually incorrect so you see this error message --- note that we got pushed into the `_init_cache_for_non_constant_diag` method for a homoscedastic noise model.
```
~/Documents/GitHub/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache(self)
    100             self._init_cache_for_constant_diag(eye, batch_shape, n, k)
    101         else:
--> 102             self._init_cache_for_non_constant_diag(eye, batch_shape, n)
    103 
    104         self._precond_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)

~/Documents/GitHub/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache_for_non_constant_diag(self, eye, batch_shape, n)
    117     def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):
    118         # With non-constant diagonals, we cant factor out the noise as easily
--> 119         self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))
    120         self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()
    121 

RuntimeError: Tensors must have same number of dimensions: got 3 and 2
```

## Expected Behavior

The noise shaping should be done properly so that `_init_cache_for_constant_diag` gets called instead. It runs properly.

## System information

**Please complete the following information:**
gpytorch 1.1.1
pytorch 1.5.0

## Additional context
I'll toss up a PR in a minute. It's not that hard to fix.



Discussion:

Resolving patch:
diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py
--- a/gpytorch/lazy/added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/added_diag_lazy_tensor.py
@@ -93,7 +93,10 @@ def precondition_closure(tensor):
     def _init_cache(self):
         *batch_shape, n, k = self._piv_chol_self.shape
         self._noise = self._diag_tensor.diag().unsqueeze(-1)
-        self._constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))
+
+        # the check for constant diag needs to be done carefully for batches.
+        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)
 
         if self._constant_diag:


Test patch: