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


--- 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 
5 import torch
6 
7 from .. import settings
8 from ..utils import broadcasting, pivoted_cholesky
9 from ..utils.warnings import NumericalWarning
10 from .diag_lazy_tensor import DiagLazyTensor
11 from .psd_sum_lazy_tensor import PsdSumLazyTensor
12 from .root_lazy_tensor import RootLazyTensor
13 from .sum_lazy_tensor import SumLazyTensor
14 
15 
16 class AddedDiagLazyTensor(SumLazyTensor):
17     """
18     A SumLazyTensor, but of only two lazy tensors, the second of which must be
19     a DiagLazyTensor.
20     """
21 
22     def __init__(self, *lazy_tensors, preconditioner_override=None):
23         lazy_tensors = list(lazy_tensors)
24         super(AddedDiagLazyTensor, self).__init__(*lazy_tensors, preconditioner_override=preconditioner_override)
25         if len(lazy_tensors) > 2:
26             raise RuntimeError("An AddedDiagLazyTensor can only have two components")
27 
28         broadcasting._mul_broadcast_shape(lazy_tensors[0].shape, lazy_tensors[1].shape)
29 
30         if isinstance(lazy_tensors[0], DiagLazyTensor) and isinstance(lazy_tensors[1], DiagLazyTensor):
31             raise RuntimeError("Trying to lazily add two DiagLazyTensors. Create a single DiagLazyTensor instead.")
32         elif isinstance(lazy_tensors[0], DiagLazyTensor):
33             self._diag_tensor = lazy_tensors[0]
34             self._lazy_tensor = lazy_tensors[1]
35         elif isinstance(lazy_tensors[1], DiagLazyTensor):
36             self._diag_tensor = lazy_tensors[1]
37             self._lazy_tensor = lazy_tensors[0]
38         else:
39             raise RuntimeError("One of the LazyTensors input to AddedDiagLazyTensor must be a DiagLazyTensor!")
40 
41         self.preconditioner_override = preconditioner_override
42 
43         # Placeholders
44         self._constant_diag = None
45         self._noise = None
46         self._piv_chol_self = None  # <- Doesn't need to be an attribute, but used for testing purposes
47         self._precond_lt = None
48         self._precond_logdet_cache = None
49         self._q_cache = None
50         self._r_cache = None
51 
52     def _matmul(self, rhs):
53         return torch.addcmul(self._lazy_tensor._matmul(rhs), self._diag_tensor._diag.unsqueeze(-1), rhs)
54 
55     def add_diag(self, added_diag):
56         return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor.add_diag(added_diag))
57 
58     def __add__(self, other):
59         from .diag_lazy_tensor import DiagLazyTensor
60 
61         if isinstance(other, DiagLazyTensor):
62             return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor + other)
63         else:
64             return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)
65 
66     def _preconditioner(self):
67         if self.preconditioner_override is not None:
68             return self.preconditioner_override(self)
69 
70         if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():
71             return None, None, None
72 
73         if self._q_cache is None:
74             max_iter = settings.max_preconditioner_size.value()
75             self._piv_chol_self = pivoted_cholesky.pivoted_cholesky(self._lazy_tensor, max_iter)
76             if torch.any(torch.isnan(self._piv_chol_self)).item():
77                 warnings.warn(
78                     "NaNs encountered in preconditioner computation. Attempting to continue without preconditioning.",
79                     NumericalWarning,
80                 )
81                 return None, None, None
82             self._init_cache()
83 
84         # NOTE: We cannot memoize this precondition closure as it causes a memory leak
85         def precondition_closure(tensor):
86             qqt = self._q_cache.matmul(self._q_cache.transpose(-2, -1).matmul(tensor))
87             if self._constant_diag:
88                 return (1 / self._noise) * (tensor - qqt)
89             return (tensor / self._noise) - qqt
90 
91         return (precondition_closure, self._precond_lt, self._precond_logdet_cache)
92 
93     def _init_cache(self):
94         *batch_shape, n, k = self._piv_chol_self.shape
95         self._noise = self._diag_tensor.diag().unsqueeze(-1)
96         self._constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))
97         eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)
98 
99         if self._constant_diag:
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)
105 
106     def _init_cache_for_constant_diag(self, eye, batch_shape, n, k):
107         # We can factor out the noise for for both QR and solves.
108         self._noise = self._noise.narrow(-2, 0, 1)
109         self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self, self._noise.sqrt() * eye), dim=-2))
110         self._q_cache = self._q_cache[..., :n, :]
111 
112         # Use the matrix determinant lemma for the logdet, using the fact that R'R = L_k'L_k + s*I
113         logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
114         logdet = logdet + (n - k) * self._noise.squeeze(-2).squeeze(-1).log()
115         self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
116 
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 
122         logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2)
123         logdet -= (1.0 / self._noise).log().sum([-1, -2])
124         self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
125
```
--- 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
@@ -93,7 +93,10 @@
     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:

{"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@@ -93,7 +93,10 @@\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-        self._constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))\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", "issue": "Batched Multi-Output GP preconditioner doesn't compute noise model properly[Bug]\n# \ud83d\udc1b Bug\r\n\r\nThe 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.\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 ExactGPModel(gpytorch.models.ExactGP):\r\n    def __init__(self, train_x, train_y, likelihood, batch_shape=torch.Size([])):\r\n        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)\r\n        self.mean_module = gpytorch.means.ConstantMean(batch_shape=batch_shape)\r\n        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(batch_shape=batch_shape),\r\n                                                        batch_shape=batch_shape)\r\n    \r\n    def forward(self, x):\r\n        mean_x = self.mean_module(x)\r\n        covar_x = self.covar_module(x)\r\n        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)\r\n\r\n## base case that works\r\ntrain_x = torch.randn(301, 1)\r\ntrain_y = torch.sin(3.1 * train_x) + 0.1 * torch.randn_like(train_x)\r\n\r\nlikelihood = gpytorch.likelihoods.GaussianLikelihood()\r\nmodel = ExactGPModel(train_x, train_y, likelihood) \r\n\r\nwith gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):\r\n\r\n    print(likelihood(model(train_x)).log_prob(train_y).sum())\r\n\r\n### This doesn't work\r\n### I'm playing with the settings here to force it to crash but it's a shaping error.\r\ntrain_x = torch.randn(301, 1)\r\ntrain_y = torch.cat((torch.sin(3.1 * train_x) + 0.1 * torch.randn_like(train_x),\r\n                     torch.sin(train_x) + 0.1 * torch.randn_like(train_x)),dim=1)\r\n\r\nlikelihood = gpytorch.likelihoods.GaussianLikelihood(batch_shape=torch.Size([2]))\r\nmodel = ExactGPModel(train_x, train_y, likelihood, batch_shape=torch.Size([2])) \r\n\r\noptimizer = torch.optim.Adam(model.parameters(), lr=0.1)\r\n\r\nfor i in range(3):\r\n    optimizer.zero_grad()\r\n\r\n    with gpytorch.settings.max_cholesky_size(100), gpytorch.settings.min_preconditioning_size(100):\r\n\r\n        loss = -likelihood(model(train_x)).log_prob(train_y.t()).sum()\r\n        loss.backward()\r\n        optimizer.step()\r\n```\r\n\r\n** Stack trace/error message **\r\nThe 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.\r\n```\r\n~/Documents/GitHub/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache(self)\r\n    100             self._init_cache_for_constant_diag(eye, batch_shape, n, k)\r\n    101         else:\r\n--> 102             self._init_cache_for_non_constant_diag(eye, batch_shape, n)\r\n    103 \r\n    104         self._precond_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)\r\n\r\n~/Documents/GitHub/gpytorch/gpytorch/lazy/added_diag_lazy_tensor.py in _init_cache_for_non_constant_diag(self, eye, batch_shape, n)\r\n    117     def _init_cache_for_non_constant_diag(self, eye, batch_shape, n):\r\n    118         # With non-constant diagonals, we cant factor out the noise as easily\r\n--> 119         self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self / self._noise.sqrt(), eye)))\r\n    120         self._q_cache = self._q_cache[..., :n, :] / self._noise.sqrt()\r\n    121 \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\nThe noise shaping should be done properly so that `_init_cache_for_constant_diag` gets called instead. It runs properly.\r\n\r\n## System information\r\n\r\n**Please complete the following information:**\r\ngpytorch 1.1.1\r\npytorch 1.5.0\r\n\r\n## Additional context\r\nI'll toss up a PR in a minute. It's not that hard to fix.\r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\n\nimport torch\n\nfrom .. import settings\nfrom ..utils import broadcasting, pivoted_cholesky\nfrom ..utils.warnings import NumericalWarning\nfrom .diag_lazy_tensor import DiagLazyTensor\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        self._constant_diag = torch.equal(self._noise, self._noise[0] * 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", "path": "gpytorch/lazy/added_diag_lazy_tensor.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\n\nimport torch\n\nfrom .. import settings\nfrom ..utils import broadcasting, pivoted_cholesky\nfrom ..utils.warnings import NumericalWarning\nfrom .diag_lazy_tensor import DiagLazyTensor\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", "path": "gpytorch/lazy/added_diag_lazy_tensor.py"}]}