We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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))
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `gpytorch/lazy/added_diag_lazy_tensor.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 import warnings
4 from typing import Optional, Tuple
5 
6 import torch
7 from torch import Tensor
8 
9 from .. import settings
10 from ..utils import broadcasting, pivoted_cholesky
11 from ..utils.memoize import cached
12 from ..utils.warnings import NumericalWarning
13 from .diag_lazy_tensor import ConstantDiagLazyTensor, DiagLazyTensor
14 from .lazy_tensor import LazyTensor
15 from .psd_sum_lazy_tensor import PsdSumLazyTensor
16 from .root_lazy_tensor import RootLazyTensor
17 from .sum_lazy_tensor import SumLazyTensor
18 
19 
20 class AddedDiagLazyTensor(SumLazyTensor):
21     """
22     A SumLazyTensor, but of only two lazy tensors, the second of which must be
23     a DiagLazyTensor.
24     """
25 
26     def __init__(self, *lazy_tensors, preconditioner_override=None):
27         lazy_tensors = list(lazy_tensors)
28         super(AddedDiagLazyTensor, self).__init__(*lazy_tensors, preconditioner_override=preconditioner_override)
29         if len(lazy_tensors) > 2:
30             raise RuntimeError("An AddedDiagLazyTensor can only have two components")
31 
32         broadcasting._mul_broadcast_shape(lazy_tensors[0].shape, lazy_tensors[1].shape)
33 
34         if isinstance(lazy_tensors[0], DiagLazyTensor) and isinstance(lazy_tensors[1], DiagLazyTensor):
35             raise RuntimeError("Trying to lazily add two DiagLazyTensors. Create a single DiagLazyTensor instead.")
36         elif isinstance(lazy_tensors[0], DiagLazyTensor):
37             self._diag_tensor = lazy_tensors[0]
38             self._lazy_tensor = lazy_tensors[1]
39         elif isinstance(lazy_tensors[1], DiagLazyTensor):
40             self._diag_tensor = lazy_tensors[1]
41             self._lazy_tensor = lazy_tensors[0]
42         else:
43             raise RuntimeError("One of the LazyTensors input to AddedDiagLazyTensor must be a DiagLazyTensor!")
44 
45         self.preconditioner_override = preconditioner_override
46 
47         # Placeholders
48         self._constant_diag = None
49         self._noise = None
50         self._piv_chol_self = None  # <- Doesn't need to be an attribute, but used for testing purposes
51         self._precond_lt = None
52         self._precond_logdet_cache = None
53         self._q_cache = None
54         self._r_cache = None
55 
56     def _matmul(self, rhs):
57         return torch.addcmul(self._lazy_tensor._matmul(rhs), self._diag_tensor._diag.unsqueeze(-1), rhs)
58 
59     def add_diag(self, added_diag):
60         return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor.add_diag(added_diag))
61 
62     def __add__(self, other):
63         from .diag_lazy_tensor import DiagLazyTensor
64 
65         if isinstance(other, DiagLazyTensor):
66             return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor + other)
67         else:
68             return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)
69 
70     def _preconditioner(self):
71         if self.preconditioner_override is not None:
72             return self.preconditioner_override(self)
73 
74         if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():
75             return None, None, None
76 
77         if self._q_cache is None:
78             max_iter = settings.max_preconditioner_size.value()
79             self._piv_chol_self = pivoted_cholesky.pivoted_cholesky(self._lazy_tensor, max_iter)
80             if torch.any(torch.isnan(self._piv_chol_self)).item():
81                 warnings.warn(
82                     "NaNs encountered in preconditioner computation. Attempting to continue without preconditioning.",
83                     NumericalWarning,
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
89         def precondition_closure(tensor):
90             qqt = self._q_cache.matmul(self._q_cache.transpose(-2, -1).matmul(tensor))
91             if self._constant_diag:
92                 return (1 / self._noise) * (tensor - qqt)
93             return (tensor / self._noise) - qqt
94 
95         return (precondition_closure, self._precond_lt, self._precond_logdet_cache)
96 
97     def _init_cache(self):
98         *batch_shape, n, k = self._piv_chol_self.shape
99         self._noise = self._diag_tensor.diag().unsqueeze(-1)
100 
101         # the check for constant diag needs to be done carefully for batches.
102         noise_first_element = self._noise[..., :1, :]
103         self._constant_diag = torch.equal(self._noise, noise_first_element * torch.ones_like(self._noise))
104         eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)
105 
106         if self._constant_diag:
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)
112 
113     def _init_cache_for_constant_diag(self, eye, batch_shape, n, k):
114         # We can factor out the noise for for both QR and solves.
115         self._noise = self._noise.narrow(-2, 0, 1)
116         self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self, self._noise.sqrt() * eye), dim=-2))
117         self._q_cache = self._q_cache[..., :n, :]
118 
119         # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
120         logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
121         logdet = logdet + (n - k) * self._noise.squeeze(-2).squeeze(-1).log()
122         self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
123 
124     def _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         self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))
127         self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()
128 
129         logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
130         logdet -= (1.0 / self._noise).log().sum([-1, -2])
131         self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
132 
133     @cached(name="svd")
134     def _svd(self) -> Tuple["LazyTensor", Tensor, "LazyTensor"]:
135         if isinstance(self._diag_tensor, ConstantDiagLazyTensor):
136             U, S_, V = self._lazy_tensor.svd()
137             S = S_ + self._diag_tensor.diag()
138             return U, S, V
139         return super()._svd()
140 
141     def _symeig(self, eigenvectors: bool = False) -> Tuple[Tensor, Optional[LazyTensor]]:
142         if isinstance(self._diag_tensor, ConstantDiagLazyTensor):
143             evals_, evecs = self._lazy_tensor.symeig(eigenvectors=eigenvectors)
144             evals = evals_ + self._diag_tensor.diag()
145             return evals, evecs
146         return super()._symeig(eigenvectors=eigenvectors)
147
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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
@@ -68,12 +68,33 @@
             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 @@
 
         # 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 @@
         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 @@
 
     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()

{"golden_diff": "diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py\n--- a/gpytorch/lazy/added_diag_lazy_tensor.py\n+++ b/gpytorch/lazy/added_diag_lazy_tensor.py\n@@ -68,12 +68,33 @@\n             return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)\n \n     def _preconditioner(self):\n+        r\"\"\"\n+        Here we use a partial pivoted Cholesky preconditioner:\n+\n+        K \\approx L L^T + D\n+\n+        where L L^T is a low rank approximation, and D is a diagonal.\n+        We can compute the preconditioner's inverse using Woodbury\n+\n+        (L L^T + D)^{-1} = D^{-1} - D^{-1} L (I + L D^{-1} L^T)^{-1} L^T D^{-1}\n+\n+        This function returns:\n+        - A function `precondition_closure` that computes the solve (L L^T + D)^{-1} x\n+        - A LazyTensor `precondition_lt` that represents (L L^T + D)\n+        - The log determinant of (L L^T + D)\n+        \"\"\"\n+\n         if self.preconditioner_override is not None:\n             return self.preconditioner_override(self)\n \n         if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():\n             return None, None, None\n \n+        # Cache a QR decomposition [Q; Q'] R = [D^{-1/2}; L]\n+        # This makes it fast to compute solves and log determinants with it\n+        #\n+        # Through woodbury, (L L^T + D)^{-1} reduces down to (D^{-1} - D^{-1/2} Q Q^T D^{-1/2})\n+        # Through matrix determinant lemma, log |L L^T + D| reduces down to 2 log |R|\n         if self._q_cache is None:\n             max_iter = settings.max_preconditioner_size.value()\n             self._piv_chol_self = pivoted_cholesky.pivoted_cholesky(self._lazy_tensor, max_iter)\n@@ -87,6 +108,7 @@\n \n         # NOTE: We cannot memoize this precondition closure as it causes a memory leak\n         def precondition_closure(tensor):\n+            # This makes it fast to compute solves with it\n             qqt = self._q_cache.matmul(self._q_cache.transpose(-2, -1).matmul(tensor))\n             if self._constant_diag:\n                 return (1 / self._noise) * (tensor - qqt)\n@@ -102,6 +124,7 @@\n         noise_first_element = self._noise[..., :1, :]\n         self._constant_diag = torch.equal(self._noise, noise_first_element * torch.ones_like(self._noise))\n         eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)\n+        eye = eye.expand(*batch_shape, k, k)\n \n         if self._constant_diag:\n             self._init_cache_for_constant_diag(eye, batch_shape, n, k)\n@@ -123,9 +146,10 @@\n \n     def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):\n         # With non-constant diagonals, we cant factor out the noise as easily\n-        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))\n+        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye), dim=-2))\n         self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()\n \n+        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I\n         logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)\n         logdet -= (1.0 / self._noise).log().sum([-1, -2])\n         self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()\n", "issue": "[Bug]\n# \ud83d\udc1b Bug\r\n\r\nBatched fixed noise GPs fail when the preconditioning threshold is reached. The concatenation in this line fails with a shape error.\r\n\r\nhttps://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/lazy/added_diag_lazy_tensor.py#L126\r\n\r\n## To reproduce\r\n\r\n** Code snippet to reproduce **\r\n```python\r\nimport torch\r\nimport gpytorch\r\n\r\nclass BatchFixedNoiseGP(gpytorch.models.GP):\r\n    def __init__(self, init_x, init_y, noise, batch_shape):\r\n        super().__init__()\r\n        self.mean_module = gpytorch.means.ZeroMean()\r\n        self.covar_module = gpytorch.kernels.RBFKernel(batch_shape=batch_shape)\r\n        self.likelihood = gpytorch.likelihoods.FixedNoiseGaussianLikelihood(noise)\r\n        \r\n    def forward(self, inputs):\r\n        mean = self.mean_module(inputs)\r\n        covar = self.covar_module(inputs)\r\n        return gpytorch.distributions.MultivariateNormal(mean, covar)\r\n\r\nbatch_shape = [2]\r\ntrain_x = torch.randn(*batch_shape, 101, 3)\r\ntrain_y = torch.randn(*batch_shape, 101)\r\ntrain_noise = torch.rand(*batch_shape, 101)\r\n\r\ngp = BatchFixedNoiseGP(train_x, train_y, train_noise, batch_shape)\r\nmll = gpytorch.mlls.ExactMarginalLogLikelihood(gp.likelihood, gp)\r\n\r\nwith gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):\r\n    train_dist = gp(train_x)\r\n    loss = -mll(train_dist, train_y).sum()\r\n```\r\n\r\n** Stack trace/error message **\r\n```\r\nRuntimeError                              Traceback (most recent call last)\r\n<ipython-input-4-9e151e2de37a> in <module>\r\n     24 with gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):\r\n     25     train_dist = gp(train_x)\r\n---> 26     loss = -mll(train_dist, train_y).sum()\r\n\r\n~/Code/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)\r\n     26 \r\n     27     def __call__(self, *inputs, **kwargs):\r\n---> 28         outputs = self.forward(*inputs, **kwargs)\r\n     29         if isinstance(outputs, list):\r\n     30             return [_validate_module_outputs(output) for output in outputs]\r\n\r\n~/Code/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, function_dist, target, *params)\r\n     49         # Get the log prob of the marginal distribution\r\n     50         output = self.likelihood(function_dist, *params)\r\n---> 51         res = output.log_prob(target)\r\n     52 \r\n     53         # Add additional terms (SGPR / learned inducing points, heteroskedastic likelihood models)\r\n\r\n~/Code/gpytorch/gpytorch/distributions/multivariate_normal.py in log_prob(self, value)\r\n    138 \r\n    139         # Get log determininat and first part of quadratic form\r\n--> 140         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)\r\n    141 \r\n    142         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])\r\n\r\n~/Code/gpytorch/gpytorch/lazy/lazy_tensor.py in inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)\r\n   1069             probe_vectors,\r\n   1070             probe_vector_norms,\r\n-> 1071             *args,\r\n   1072         )\r\n   1073 \r\n\r\n~/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)\r\n     65         lazy_tsr = ctx.representation_tree(*matrix_args)\r\n     66         with torch.no_grad():\r\n---> 67             preconditioner, precond_lt, logdet_correction = lazy_tsr._preconditioner()\r\n     68 \r\n     69         ctx.preconditioner = preconditioner\r\n\r\n~/Code/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _preconditioner(self)\r\n     84                 )\r\n     85                 return None, None, None\r\n---> 86             self._init_cache()\r\n     87 \r\n     88         # NOTE: We cannot memoize this precondition closure as it causes a memory leak\r\n\r\n~/Code/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache(self)\r\n    107             self._init_cache_for_constant_diag(eye, batch_shape, n, k)\r\n    108         else:\r\n--> 109             self._init_cache_for_non_constant_diag(eye, batch_shape, n)\r\n    110 \r\n    111         self._precond_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)\r\n\r\n~/Code/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache_for_non_constant_diag(self, eye, batch_shape, n)\r\n    125         # With non-constant diagonals, we cant factor out the noise as easily\r\n    126         # eye = eye.expand(*batch_shape, -1, -1)\r\n--> 127         self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))\r\n    128         self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()\r\n    129 \r\n\r\nRuntimeError: Tensors must have same number of dimensions: got 3 and 2\r\n```\r\n\r\n## Expected Behavior\r\n\r\nEverything works fine until the preconditioning threshold is reached. Obviously one would hope that it would continue to work.\r\n\r\n## System information\r\n\r\n**Please complete the following information:**\r\nGPyTorch Version: 1.2.0\r\nPyTorch Version:  '1.6.0.dev20200522'\r\nOS: Ubuntu 16.04 LTS\r\n\r\n## Additional context\r\nThis appears to fix the problem\r\n```\r\n    def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):\r\n        # With non-constant diagonals, we cant factor out the noise as easily\r\n        eye = eye.expand(*batch_shape, -1, -1)\r\n        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye), dim=-2))\r\n```\r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\nfrom typing import Optional, Tuple\n\nimport torch\nfrom torch import Tensor\n\nfrom .. import settings\nfrom ..utils import broadcasting, pivoted_cholesky\nfrom ..utils.memoize import cached\nfrom ..utils.warnings import NumericalWarning\nfrom .diag_lazy_tensor import ConstantDiagLazyTensor, DiagLazyTensor\nfrom .lazy_tensor import LazyTensor\nfrom .psd_sum_lazy_tensor import PsdSumLazyTensor\nfrom .root_lazy_tensor import RootLazyTensor\nfrom .sum_lazy_tensor import SumLazyTensor\n\n\nclass AddedDiagLazyTensor(SumLazyTensor):\n    \"\"\"\n    A SumLazyTensor, but of only two lazy tensors, the second of which must be\n    a DiagLazyTensor.\n    \"\"\"\n\n    def __init__(self, *lazy_tensors, preconditioner_override=None):\n        lazy_tensors = list(lazy_tensors)\n        super(AddedDiagLazyTensor, self).__init__(*lazy_tensors, preconditioner_override=preconditioner_override)\n        if len(lazy_tensors) > 2:\n            raise RuntimeError(\"An AddedDiagLazyTensor can only have two components\")\n\n        broadcasting._mul_broadcast_shape(lazy_tensors[0].shape, lazy_tensors[1].shape)\n\n        if isinstance(lazy_tensors[0], DiagLazyTensor) and isinstance(lazy_tensors[1], DiagLazyTensor):\n            raise RuntimeError(\"Trying to lazily add two DiagLazyTensors. Create a single DiagLazyTensor instead.\")\n        elif isinstance(lazy_tensors[0], DiagLazyTensor):\n            self._diag_tensor = lazy_tensors[0]\n            self._lazy_tensor = lazy_tensors[1]\n        elif isinstance(lazy_tensors[1], DiagLazyTensor):\n            self._diag_tensor = lazy_tensors[1]\n            self._lazy_tensor = lazy_tensors[0]\n        else:\n            raise RuntimeError(\"One of the LazyTensors input to AddedDiagLazyTensor must be a DiagLazyTensor!\")\n\n        self.preconditioner_override = preconditioner_override\n\n        # Placeholders\n        self._constant_diag = None\n        self._noise = None\n        self._piv_chol_self = None  # <- Doesn't need to be an attribute, but used for testing purposes\n        self._precond_lt = None\n        self._precond_logdet_cache = None\n        self._q_cache = None\n        self._r_cache = None\n\n    def _matmul(self, rhs):\n        return torch.addcmul(self._lazy_tensor._matmul(rhs), self._diag_tensor._diag.unsqueeze(-1), rhs)\n\n    def add_diag(self, added_diag):\n        return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor.add_diag(added_diag))\n\n    def __add__(self, other):\n        from .diag_lazy_tensor import DiagLazyTensor\n\n        if isinstance(other, DiagLazyTensor):\n            return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor + other)\n        else:\n            return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)\n\n    def _preconditioner(self):\n        if self.preconditioner_override is not None:\n            return self.preconditioner_override(self)\n\n        if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():\n            return None, None, None\n\n        if self._q_cache is None:\n            max_iter = settings.max_preconditioner_size.value()\n            self._piv_chol_self = pivoted_cholesky.pivoted_cholesky(self._lazy_tensor, max_iter)\n            if torch.any(torch.isnan(self._piv_chol_self)).item():\n                warnings.warn(\n                    \"NaNs encountered in preconditioner computation. Attempting to continue without preconditioning.\",\n                    NumericalWarning,\n                )\n                return None, None, None\n            self._init_cache()\n\n        # NOTE: We cannot memoize this precondition closure as it causes a memory leak\n        def precondition_closure(tensor):\n            qqt = self._q_cache.matmul(self._q_cache.transpose(-2, -1).matmul(tensor))\n            if self._constant_diag:\n                return (1 / self._noise) * (tensor - qqt)\n            return (tensor / self._noise) - qqt\n\n        return (precondition_closure, self._precond_lt, self._precond_logdet_cache)\n\n    def _init_cache(self):\n        *batch_shape, n, k = self._piv_chol_self.shape\n        self._noise = self._diag_tensor.diag().unsqueeze(-1)\n\n        # the check for constant diag needs to be done carefully for batches.\n        noise_first_element = self._noise[..., :1, :]\n        self._constant_diag = torch.equal(self._noise, noise_first_element * torch.ones_like(self._noise))\n        eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)\n\n        if self._constant_diag:\n            self._init_cache_for_constant_diag(eye, batch_shape, n, k)\n        else:\n            self._init_cache_for_non_constant_diag(eye, batch_shape, n)\n\n        self._precond_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)\n\n    def _init_cache_for_constant_diag(self, eye, batch_shape, n, k):\n        # We can factor out the noise for for both QR and solves.\n        self._noise = self._noise.narrow(-2, 0, 1)\n        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self, self._noise.sqrt() * eye), dim=-2))\n        self._q_cache = self._q_cache[..., :n, :]\n\n        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I\n        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)\n        logdet = logdet + (n - k) * self._noise.squeeze(-2).squeeze(-1).log()\n        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()\n\n    def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):\n        # With non-constant diagonals, we cant factor out the noise as easily\n        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))\n        self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()\n\n        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)\n        logdet -= (1.0 / self._noise).log().sum([-1, -2])\n        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()\n\n    @cached(name=\"svd\")\n    def _svd(self) -> Tuple[\"LazyTensor\", Tensor, \"LazyTensor\"]:\n        if isinstance(self._diag_tensor, ConstantDiagLazyTensor):\n            U, S_, V = self._lazy_tensor.svd()\n            S = S_ + self._diag_tensor.diag()\n            return U, S, V\n        return super()._svd()\n\n    def _symeig(self, eigenvectors: bool = False) -> Tuple[Tensor, Optional[LazyTensor]]:\n        if isinstance(self._diag_tensor, ConstantDiagLazyTensor):\n            evals_, evecs = self._lazy_tensor.symeig(eigenvectors=eigenvectors)\n            evals = evals_ + self._diag_tensor.diag()\n            return evals, evecs\n        return super()._symeig(eigenvectors=eigenvectors)\n", "path": "gpytorch/lazy/added_diag_lazy_tensor.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\nfrom typing import Optional, Tuple\n\nimport torch\nfrom torch import Tensor\n\nfrom .. import settings\nfrom ..utils import broadcasting, pivoted_cholesky\nfrom ..utils.memoize import cached\nfrom ..utils.warnings import NumericalWarning\nfrom .diag_lazy_tensor import ConstantDiagLazyTensor, DiagLazyTensor\nfrom .lazy_tensor import LazyTensor\nfrom .psd_sum_lazy_tensor import PsdSumLazyTensor\nfrom .root_lazy_tensor import RootLazyTensor\nfrom .sum_lazy_tensor import SumLazyTensor\n\n\nclass AddedDiagLazyTensor(SumLazyTensor):\n    \"\"\"\n    A SumLazyTensor, but of only two lazy tensors, the second of which must be\n    a DiagLazyTensor.\n    \"\"\"\n\n    def __init__(self, *lazy_tensors, preconditioner_override=None):\n        lazy_tensors = list(lazy_tensors)\n        super(AddedDiagLazyTensor, self).__init__(*lazy_tensors, preconditioner_override=preconditioner_override)\n        if len(lazy_tensors) > 2:\n            raise RuntimeError(\"An AddedDiagLazyTensor can only have two components\")\n\n        broadcasting._mul_broadcast_shape(lazy_tensors[0].shape, lazy_tensors[1].shape)\n\n        if isinstance(lazy_tensors[0], DiagLazyTensor) and isinstance(lazy_tensors[1], DiagLazyTensor):\n            raise RuntimeError(\"Trying to lazily add two DiagLazyTensors. Create a single DiagLazyTensor instead.\")\n        elif isinstance(lazy_tensors[0], DiagLazyTensor):\n            self._diag_tensor = lazy_tensors[0]\n            self._lazy_tensor = lazy_tensors[1]\n        elif isinstance(lazy_tensors[1], DiagLazyTensor):\n            self._diag_tensor = lazy_tensors[1]\n            self._lazy_tensor = lazy_tensors[0]\n        else:\n            raise RuntimeError(\"One of the LazyTensors input to AddedDiagLazyTensor must be a DiagLazyTensor!\")\n\n        self.preconditioner_override = preconditioner_override\n\n        # Placeholders\n        self._constant_diag = None\n        self._noise = None\n        self._piv_chol_self = None  # <- Doesn't need to be an attribute, but used for testing purposes\n        self._precond_lt = None\n        self._precond_logdet_cache = None\n        self._q_cache = None\n        self._r_cache = None\n\n    def _matmul(self, rhs):\n        return torch.addcmul(self._lazy_tensor._matmul(rhs), self._diag_tensor._diag.unsqueeze(-1), rhs)\n\n    def add_diag(self, added_diag):\n        return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor.add_diag(added_diag))\n\n    def __add__(self, other):\n        from .diag_lazy_tensor import DiagLazyTensor\n\n        if isinstance(other, DiagLazyTensor):\n            return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor + other)\n        else:\n            return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)\n\n    def _preconditioner(self):\n        r\"\"\"\n        Here we use a partial pivoted Cholesky preconditioner:\n\n        K \\approx L L^T + D\n\n        where L L^T is a low rank approximation, and D is a diagonal.\n        We can compute the preconditioner's inverse using Woodbury\n\n        (L L^T + D)^{-1} = D^{-1} - D^{-1} L (I + L D^{-1} L^T)^{-1} L^T D^{-1}\n\n        This function returns:\n        - A function `precondition_closure` that computes the solve (L L^T + D)^{-1} x\n        - A LazyTensor `precondition_lt` that represents (L L^T + D)\n        - The log determinant of (L L^T + D)\n        \"\"\"\n\n        if self.preconditioner_override is not None:\n            return self.preconditioner_override(self)\n\n        if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():\n            return None, None, None\n\n        # Cache a QR decomposition [Q; Q'] R = [D^{-1/2}; L]\n        # This makes it fast to compute solves and log determinants with it\n        #\n        # Through woodbury, (L L^T + D)^{-1} reduces down to (D^{-1} - D^{-1/2} Q Q^T D^{-1/2})\n        # Through matrix determinant lemma, log |L L^T + D| reduces down to 2 log |R|\n        if self._q_cache is None:\n            max_iter = settings.max_preconditioner_size.value()\n            self._piv_chol_self = pivoted_cholesky.pivoted_cholesky(self._lazy_tensor, max_iter)\n            if torch.any(torch.isnan(self._piv_chol_self)).item():\n                warnings.warn(\n                    \"NaNs encountered in preconditioner computation. Attempting to continue without preconditioning.\",\n                    NumericalWarning,\n                )\n                return None, None, None\n            self._init_cache()\n\n        # NOTE: We cannot memoize this precondition closure as it causes a memory leak\n        def precondition_closure(tensor):\n            # This makes it fast to compute solves with it\n            qqt = self._q_cache.matmul(self._q_cache.transpose(-2, -1).matmul(tensor))\n            if self._constant_diag:\n                return (1 / self._noise) * (tensor - qqt)\n            return (tensor / self._noise) - qqt\n\n        return (precondition_closure, self._precond_lt, self._precond_logdet_cache)\n\n    def _init_cache(self):\n        *batch_shape, n, k = self._piv_chol_self.shape\n        self._noise = self._diag_tensor.diag().unsqueeze(-1)\n\n        # the check for constant diag needs to be done carefully for batches.\n        noise_first_element = self._noise[..., :1, :]\n        self._constant_diag = torch.equal(self._noise, noise_first_element * torch.ones_like(self._noise))\n        eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)\n        eye = eye.expand(*batch_shape, k, k)\n\n        if self._constant_diag:\n            self._init_cache_for_constant_diag(eye, batch_shape, n, k)\n        else:\n            self._init_cache_for_non_constant_diag(eye, batch_shape, n)\n\n        self._precond_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)\n\n    def _init_cache_for_constant_diag(self, eye, batch_shape, n, k):\n        # We can factor out the noise for for both QR and solves.\n        self._noise = self._noise.narrow(-2, 0, 1)\n        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self, self._noise.sqrt() * eye), dim=-2))\n        self._q_cache = self._q_cache[..., :n, :]\n\n        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I\n        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)\n        logdet = logdet + (n - k) * self._noise.squeeze(-2).squeeze(-1).log()\n        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()\n\n    def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):\n        # With non-constant diagonals, we cant factor out the noise as easily\n        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye), dim=-2))\n        self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()\n\n        # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I\n        logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)\n        logdet -= (1.0 / self._noise).log().sum([-1, -2])\n        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()\n\n    @cached(name=\"svd\")\n    def _svd(self) -> Tuple[\"LazyTensor\", Tensor, \"LazyTensor\"]:\n        if isinstance(self._diag_tensor, ConstantDiagLazyTensor):\n            U, S_, V = self._lazy_tensor.svd()\n            S = S_ + self._diag_tensor.diag()\n            return U, S, V\n        return super()._svd()\n\n    def _symeig(self, eigenvectors: bool = False) -> Tuple[Tensor, Optional[LazyTensor]]:\n        if isinstance(self._diag_tensor, ConstantDiagLazyTensor):\n            evals_, evecs = self._lazy_tensor.symeig(eigenvectors=eigenvectors)\n            evals = evals_ + self._diag_tensor.diag()\n            return evals, evecs\n        return super()._symeig(eigenvectors=eigenvectors)\n", "path": "gpytorch/lazy/added_diag_lazy_tensor.py"}]}