We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Bug] `gpytorch.settings.deterministic_probes` computes an extra Lanczos factorization
# 🐛 Bug

In preconditioned CG mode, we make the following lining calls for each training iteration:

```
LinAlg (Verbose) - DEBUG - Running Pivoted Cholesky on a torch.Size([5000, 5000]) RHS for 15 iterations.

# these two calls shouldn't exist
LinAlg (Verbose) - DEBUG - Running Lanczos on a torch.Size([5000, 5000]) matrix with a torch.Size([5000, 1]) RHS for 100 iterations.
LinAlg (Verbose) - DEBUG - Running symeig on a matrix of size torch.Size([9, 9]).

LinAlg (Verbose) - DEBUG - Running CG on a torch.Size([5000, 11]) RHS for 1000 iterations (tol=1). Output: torch.Size([5000, 11]).
LinAlg (Verbose) - DEBUG - Running symeig on a matrix of size torch.Size([10, 4, 4]).

LinAlg (Verbose) - INFO - Iter 1/50 - Loss: 0.808   lengthscale: 0.693   noise: 0.693
```

This Lanczos/symeig calls come from this line in [inv_quad_logdet](https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/functions/_inv_quad_log_det.py#L98)

```python
            else:  # When preconditioning, probe vectors must be drawn from N(0, P)
                if precond_lt.size()[-2:] == torch.Size([1, 1]):
                    covar_root = precond_lt.evaluate().sqrt()
                else:
                    covar_root = precond_lt.root_decomposition().root
```

Before the `gpytorch.settings.deterministic_probes` option was added (#929), the code looked like this:

```python
            else:  # When preconditioning, probe vectors must be drawn from N(0, P)
                probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)
                # ...
```

This is much more efficient, as it does not perform a root decomposition.

@jacobrgardner there is an obvious easy fix: we don't call `root_decomposition` unless we're using deterministic probes. However, I think this makes the `deterministic_probes` setting really undesirable - the slight variance reductions gains we get from it come at the cost of an extra linalg routine.

Thoughts?

--- 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/functions/_inv_quad_log_det.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 import warnings
4 
5 import torch
6 from torch.autograd import Function
7 
8 from .. import settings
9 from ..utils.lanczos import lanczos_tridiag_to_diag
10 from ..utils.stochastic_lq import StochasticLQ
11 
12 
13 class InvQuadLogDet(Function):
14     """
15     Given a PSD matrix A (or a batch of PSD matrices A), this function computes one or both
16     of the following
17     - The matrix solves A^{-1} b
18     - logdet(A)
19     """
20 
21     @staticmethod
22     def forward(
23         ctx,
24         representation_tree,
25         dtype,
26         device,
27         matrix_shape,
28         batch_shape=torch.Size(),
29         inv_quad=False,
30         logdet=False,
31         probe_vectors=None,
32         probe_vector_norms=None,
33         *args,
34     ):
35         """
36         *args - The arguments representing the PSD matrix A (or batch of PSD matrices A)
37         If self.inv_quad is true, the first entry in *args is inv_quad_rhs (Tensor)
38         - the RHS of the matrix solves.
39 
40         Returns:
41         - (Scalar) The inverse quadratic form (or None, if self.inv_quad is False)
42         - (Scalar) The log determinant (or None, self.if logdet is False)
43         """
44 
45         if not (inv_quad or logdet):
46             raise RuntimeError("Either inv_quad or logdet must be true (or both)")
47 
48         ctx.representation_tree = representation_tree
49         ctx.dtype = dtype
50         ctx.device = device
51         ctx.matrix_shape = matrix_shape
52         ctx.batch_shape = batch_shape
53         ctx.inv_quad = inv_quad
54         ctx.logdet = logdet
55 
56         matrix_args = None
57         inv_quad_rhs = None
58         if ctx.inv_quad:
59             matrix_args = args[1:]
60             inv_quad_rhs = args[0]
61         else:
62             matrix_args = args
63 
64         # Get closure for matmul
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
70 
71         if (probe_vectors is None or probe_vector_norms is None) and logdet:
72             num_random_probes = settings.num_trace_samples.value()
73             if preconditioner is None:
74                 if settings.deterministic_probes.on():
75                     warnings.warn(
76                         "Deterministic probes will currently work only if you aren't training multiple independent"
77                         " models simultaneously.",
78                         UserWarning,
79                     )
80                     if settings.deterministic_probes.probe_vectors is None:
81                         probe_vectors = torch.empty(matrix_shape[-1], num_random_probes, dtype=dtype, device=device)
82                         probe_vectors.bernoulli_().mul_(2).add_(-1)
83                         settings.deterministic_probes.probe_vectors = probe_vectors
84                     else:
85                         probe_vectors = settings.deterministic_probes.probe_vectors
86                 else:
87                     probe_vectors = torch.empty(matrix_shape[-1], num_random_probes, dtype=dtype, device=device)
88                     probe_vectors.bernoulli_().mul_(2).add_(-1)
89 
90                 probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)
91                 if batch_shape is not None:
92                     probe_vectors = probe_vectors.expand(*batch_shape, matrix_shape[-1], num_random_probes)
93                     probe_vector_norms = probe_vector_norms.expand(*batch_shape, 1, num_random_probes)
94             else:  # When preconditioning, probe vectors must be drawn from N(0, P)
95                 if precond_lt.size()[-2:] == torch.Size([1, 1]):
96                     covar_root = precond_lt.evaluate().sqrt()
97                 else:
98                     covar_root = precond_lt.root_decomposition().root
99 
100                 if settings.deterministic_probes.on():
101                     warnings.warn(
102                         "Deterministic probes will currently work only if you aren't training multiple independent"
103                         " models simultaneously.",
104                         UserWarning,
105                     )
106                     base_samples = settings.deterministic_probes.probe_vectors
107                     if base_samples is None or covar_root.size(-1) != base_samples.size(-2):
108                         base_samples = torch.randn(
109                             *precond_lt.batch_shape,
110                             covar_root.size(-1),
111                             num_random_probes,
112                             dtype=precond_lt.dtype,
113                             device=precond_lt.device,
114                         )
115                         settings.deterministic_probes.probe_vectors = base_samples
116 
117                     probe_vectors = covar_root.matmul(base_samples).permute(-1, *range(precond_lt.dim() - 1))
118                 else:
119                     base_samples = torch.randn(
120                         *precond_lt.batch_shape,
121                         covar_root.size(-1),
122                         num_random_probes,
123                         dtype=precond_lt.dtype,
124                         device=precond_lt.device,
125                     )
126                     probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)
127                 probe_vectors = probe_vectors.unsqueeze(-2).transpose(0, -2).squeeze(0).transpose(-2, -1).contiguous()
128                 probe_vector_norms = torch.norm(probe_vectors, p=2, dim=-2, keepdim=True)
129             probe_vectors = probe_vectors.div(probe_vector_norms)
130 
131         ctx.probe_vectors = probe_vectors
132         ctx.probe_vector_norms = probe_vector_norms
133 
134         if ctx.logdet and not ctx.probe_vectors.numel():
135             raise RuntimeError("Probe vectors were not supplied for logdet computation")
136 
137         # Collect terms for LinearCG
138         # We use LinearCG for both matrix solves and for stochastically estimating the log det
139         rhs_list = []
140         num_random_probes = 0
141         num_inv_quad_solves = 0
142 
143         # RHS for logdet
144         if ctx.logdet:
145             rhs_list.append(ctx.probe_vectors)
146             num_random_probes = ctx.probe_vectors.size(-1)
147 
148         # RHS for inv_quad
149         ctx.is_vector = False
150         if ctx.inv_quad:
151             if inv_quad_rhs.ndimension() == 1:
152                 inv_quad_rhs = inv_quad_rhs.unsqueeze(-1)
153                 ctx.is_vector = True
154             rhs_list.append(inv_quad_rhs)
155             num_inv_quad_solves = inv_quad_rhs.size(-1)
156 
157         # Perform solves (for inv_quad) and tridiagonalization (for estimating logdet)
158         rhs = torch.cat(rhs_list, -1)
159         t_mat = None
160         if ctx.logdet and settings.skip_logdet_forward.off():
161             solves, t_mat = lazy_tsr._solve(rhs, preconditioner, num_tridiag=num_random_probes)
162 
163         else:
164             solves = lazy_tsr._solve(rhs, preconditioner, num_tridiag=0)
165 
166         # Final values to return
167         logdet_term = torch.zeros(lazy_tsr.batch_shape, dtype=ctx.dtype, device=ctx.device)
168         inv_quad_term = torch.zeros(lazy_tsr.batch_shape, dtype=ctx.dtype, device=ctx.device)
169 
170         # Compute logdet from tridiagonalization
171         if ctx.logdet and settings.skip_logdet_forward.off():
172             if torch.any(torch.isnan(t_mat)).item():
173                 logdet_term = torch.tensor(float("nan"), dtype=ctx.dtype, device=ctx.device)
174             else:
175                 if ctx.batch_shape is None:
176                     t_mat = t_mat.unsqueeze(1)
177                 eigenvalues, eigenvectors = lanczos_tridiag_to_diag(t_mat)
178                 slq = StochasticLQ()
179                 (logdet_term,) = slq.evaluate(ctx.matrix_shape, eigenvalues, eigenvectors, [lambda x: x.log()])
180 
181                 # Add correction
182                 if logdet_correction is not None:
183                     logdet_term = logdet_term + logdet_correction
184 
185         # Extract inv_quad solves from all the solves
186         if ctx.inv_quad:
187             inv_quad_solves = solves.narrow(-1, num_random_probes, num_inv_quad_solves)
188             inv_quad_term = (inv_quad_solves * inv_quad_rhs).sum(-2)
189 
190         ctx.num_random_probes = num_random_probes
191         ctx.num_inv_quad_solves = num_inv_quad_solves
192 
193         to_save = list(matrix_args) + [solves]
194         ctx.save_for_backward(*to_save)
195 
196         if settings.memory_efficient.off():
197             ctx._lazy_tsr = lazy_tsr
198 
199         return inv_quad_term, logdet_term
200 
201     @staticmethod
202     def backward(ctx, inv_quad_grad_output, logdet_grad_output):
203         matrix_arg_grads = None
204         inv_quad_rhs_grad = None
205 
206         # Which backward passes should we compute?
207         compute_inv_quad_grad = inv_quad_grad_output.abs().sum() and ctx.inv_quad
208         compute_logdet_grad = logdet_grad_output.abs().sum() and ctx.logdet
209 
210         # Get input arguments, and get gradients in the proper form
211         matrix_args = ctx.saved_tensors[:-1]
212         solves = ctx.saved_tensors[-1]
213 
214         if hasattr(ctx, "_lazy_tsr"):
215             lazy_tsr = ctx._lazy_tsr
216         else:
217             lazy_tsr = ctx.representation_tree(*matrix_args)
218 
219         # Fix grad_output sizes
220         if ctx.inv_quad:
221             inv_quad_grad_output = inv_quad_grad_output.unsqueeze(-2)
222         if compute_logdet_grad:
223             logdet_grad_output = logdet_grad_output.unsqueeze(-1)
224             logdet_grad_output.unsqueeze_(-1)
225 
226         # Divide up the solves
227         probe_vector_solves = None
228         inv_quad_solves = None
229         neg_inv_quad_solves_times_grad_out = None
230         if compute_logdet_grad:
231             coef = 1.0 / ctx.probe_vectors.size(-1)
232             probe_vector_solves = solves.narrow(-1, 0, ctx.num_random_probes).mul(coef)
233             probe_vector_solves.mul_(ctx.probe_vector_norms).mul_(logdet_grad_output)
234             probe_vectors = ctx.probe_vectors.mul(ctx.probe_vector_norms)
235         if ctx.inv_quad:
236             inv_quad_solves = solves.narrow(-1, ctx.num_random_probes, ctx.num_inv_quad_solves)
237             neg_inv_quad_solves_times_grad_out = inv_quad_solves.mul(inv_quad_grad_output).mul_(-1)
238 
239         # input_1 gradient
240         if any(ctx.needs_input_grad):
241             # Collect terms for arg grads
242             left_factors_list = []
243             right_factors_list = []
244 
245             if compute_logdet_grad:
246                 left_factors_list.append(probe_vector_solves)
247                 if ctx.preconditioner is not None:
248                     probe_vectors = ctx.preconditioner(probe_vectors)
249                 right_factors_list.append(probe_vectors)
250 
251             if compute_inv_quad_grad:
252                 left_factors_list.append(neg_inv_quad_solves_times_grad_out)
253                 right_factors_list.append(inv_quad_solves)
254 
255             left_factors = torch.cat(left_factors_list, -1)
256             right_factors = torch.cat(right_factors_list, -1)
257             matrix_arg_grads = lazy_tsr._quad_form_derivative(left_factors, right_factors)
258 
259         # input_2 gradients
260         if compute_inv_quad_grad and ctx.needs_input_grad[9]:
261             inv_quad_rhs_grad = neg_inv_quad_solves_times_grad_out.mul_(-2)
262         elif ctx.inv_quad:
263             inv_quad_rhs_grad = torch.zeros_like(inv_quad_solves)
264         if ctx.is_vector:
265             inv_quad_rhs_grad.squeeze_(-1)
266 
267         if ctx.inv_quad:
268             res = [inv_quad_rhs_grad] + list(matrix_arg_grads)
269         else:
270             res = list(matrix_arg_grads)
271 
272         return tuple([None] * 9 + res)
273
```
--- 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/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -92,12 +92,15 @@
                     probe_vectors = probe_vectors.expand(*batch_shape, matrix_shape[-1], num_random_probes)
                     probe_vector_norms = probe_vector_norms.expand(*batch_shape, 1, num_random_probes)
             else:  # When preconditioning, probe vectors must be drawn from N(0, P)
-                if precond_lt.size()[-2:] == torch.Size([1, 1]):
-                    covar_root = precond_lt.evaluate().sqrt()
-                else:
-                    covar_root = precond_lt.root_decomposition().root
-
                 if settings.deterministic_probes.on():
+                    # NOTE: calling precond_lt.root_decomposition() is expensive
+                    # because it requires Lanczos
+                    # We don't have any other choice for when we want to use deterministic probes, however
+                    if precond_lt.size()[-2:] == torch.Size([1, 1]):
+                        covar_root = precond_lt.evaluate().sqrt()
+                    else:
+                        covar_root = precond_lt.root_decomposition().root
+
                     warnings.warn(
                         "Deterministic probes will currently work only if you aren't training multiple independent"
                         " models simultaneously.",
@@ -116,13 +119,6 @@
 
                     probe_vectors = covar_root.matmul(base_samples).permute(-1, *range(precond_lt.dim() - 1))
                 else:
-                    base_samples = torch.randn(
-                        *precond_lt.batch_shape,
-                        covar_root.size(-1),
-                        num_random_probes,
-                        dtype=precond_lt.dtype,
-                        device=precond_lt.device,
-                    )
                     probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)
                 probe_vectors = probe_vectors.unsqueeze(-2).transpose(0, -2).squeeze(0).transpose(-2, -1).contiguous()
                 probe_vector_norms = torch.norm(probe_vectors, p=2, dim=-2, keepdim=True)

{"golden_diff": "diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py\n--- a/gpytorch/functions/_inv_quad_log_det.py\n+++ b/gpytorch/functions/_inv_quad_log_det.py\n@@ -92,12 +92,15 @@\n                     probe_vectors = probe_vectors.expand(*batch_shape, matrix_shape[-1], num_random_probes)\n                     probe_vector_norms = probe_vector_norms.expand(*batch_shape, 1, num_random_probes)\n             else:  # When preconditioning, probe vectors must be drawn from N(0, P)\n-                if precond_lt.size()[-2:] == torch.Size([1, 1]):\n-                    covar_root = precond_lt.evaluate().sqrt()\n-                else:\n-                    covar_root = precond_lt.root_decomposition().root\n-\n                 if settings.deterministic_probes.on():\n+                    # NOTE: calling precond_lt.root_decomposition() is expensive\n+                    # because it requires Lanczos\n+                    # We don't have any other choice for when we want to use deterministic probes, however\n+                    if precond_lt.size()[-2:] == torch.Size([1, 1]):\n+                        covar_root = precond_lt.evaluate().sqrt()\n+                    else:\n+                        covar_root = precond_lt.root_decomposition().root\n+\n                     warnings.warn(\n                         \"Deterministic probes will currently work only if you aren't training multiple independent\"\n                         \" models simultaneously.\",\n@@ -116,13 +119,6 @@\n \n                     probe_vectors = covar_root.matmul(base_samples).permute(-1, *range(precond_lt.dim() - 1))\n                 else:\n-                    base_samples = torch.randn(\n-                        *precond_lt.batch_shape,\n-                        covar_root.size(-1),\n-                        num_random_probes,\n-                        dtype=precond_lt.dtype,\n-                        device=precond_lt.device,\n-                    )\n                     probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)\n                 probe_vectors = probe_vectors.unsqueeze(-2).transpose(0, -2).squeeze(0).transpose(-2, -1).contiguous()\n                 probe_vector_norms = torch.norm(probe_vectors, p=2, dim=-2, keepdim=True)\n", "issue": "[Bug] `gpytorch.settings.deterministic_probes` computes an extra Lanczos factorization\n# \ud83d\udc1b Bug\r\n\r\nIn preconditioned CG mode, we make the following lining calls for each training iteration:\r\n\r\n```\r\nLinAlg (Verbose) - DEBUG - Running Pivoted Cholesky on a torch.Size([5000, 5000]) RHS for 15 iterations.\r\n\r\n# these two calls shouldn't exist\r\nLinAlg (Verbose) - DEBUG - Running Lanczos on a torch.Size([5000, 5000]) matrix with a torch.Size([5000, 1]) RHS for 100 iterations.\r\nLinAlg (Verbose) - DEBUG - Running symeig on a matrix of size torch.Size([9, 9]).\r\n\r\nLinAlg (Verbose) - DEBUG - Running CG on a torch.Size([5000, 11]) RHS for 1000 iterations (tol=1). Output: torch.Size([5000, 11]).\r\nLinAlg (Verbose) - DEBUG - Running symeig on a matrix of size torch.Size([10, 4, 4]).\r\n\r\nLinAlg (Verbose) - INFO - Iter 1/50 - Loss: 0.808   lengthscale: 0.693   noise: 0.693\r\n```\r\n\r\nThis Lanczos/symeig calls come from this line in [inv_quad_logdet](https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/functions/_inv_quad_log_det.py#L98)\r\n\r\n```python\r\n            else:  # When preconditioning, probe vectors must be drawn from N(0, P)\r\n                if precond_lt.size()[-2:] == torch.Size([1, 1]):\r\n                    covar_root = precond_lt.evaluate().sqrt()\r\n                else:\r\n                    covar_root = precond_lt.root_decomposition().root\r\n```\r\n\r\nBefore the `gpytorch.settings.deterministic_probes` option was added (#929), the code looked like this:\r\n\r\n```python\r\n            else:  # When preconditioning, probe vectors must be drawn from N(0, P)\r\n                probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)\r\n                # ...\r\n```\r\n\r\nThis is much more efficient, as it does not perform a root decomposition.\r\n\r\n@jacobrgardner there is an obvious easy fix: we don't call `root_decomposition` unless we're using deterministic probes. However, I think this makes the `deterministic_probes` setting really undesirable - the slight variance reductions gains we get from it come at the cost of an extra linalg routine.\r\n\r\nThoughts?\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\n\nimport torch\nfrom torch.autograd import Function\n\nfrom .. import settings\nfrom ..utils.lanczos import lanczos_tridiag_to_diag\nfrom ..utils.stochastic_lq import StochasticLQ\n\n\nclass InvQuadLogDet(Function):\n    \"\"\"\n    Given a PSD matrix A (or a batch of PSD matrices A), this function computes one or both\n    of the following\n    - The matrix solves A^{-1} b\n    - logdet(A)\n    \"\"\"\n\n    @staticmethod\n    def forward(\n        ctx,\n        representation_tree,\n        dtype,\n        device,\n        matrix_shape,\n        batch_shape=torch.Size(),\n        inv_quad=False,\n        logdet=False,\n        probe_vectors=None,\n        probe_vector_norms=None,\n        *args,\n    ):\n        \"\"\"\n        *args - The arguments representing the PSD matrix A (or batch of PSD matrices A)\n        If self.inv_quad is true, the first entry in *args is inv_quad_rhs (Tensor)\n        - the RHS of the matrix solves.\n\n        Returns:\n        - (Scalar) The inverse quadratic form (or None, if self.inv_quad is False)\n        - (Scalar) The log determinant (or None, self.if logdet is False)\n        \"\"\"\n\n        if not (inv_quad or logdet):\n            raise RuntimeError(\"Either inv_quad or logdet must be true (or both)\")\n\n        ctx.representation_tree = representation_tree\n        ctx.dtype = dtype\n        ctx.device = device\n        ctx.matrix_shape = matrix_shape\n        ctx.batch_shape = batch_shape\n        ctx.inv_quad = inv_quad\n        ctx.logdet = logdet\n\n        matrix_args = None\n        inv_quad_rhs = None\n        if ctx.inv_quad:\n            matrix_args = args[1:]\n            inv_quad_rhs = args[0]\n        else:\n            matrix_args = args\n\n        # Get closure for matmul\n        lazy_tsr = ctx.representation_tree(*matrix_args)\n        with torch.no_grad():\n            preconditioner, precond_lt, logdet_correction = lazy_tsr._preconditioner()\n\n        ctx.preconditioner = preconditioner\n\n        if (probe_vectors is None or probe_vector_norms is None) and logdet:\n            num_random_probes = settings.num_trace_samples.value()\n            if preconditioner is None:\n                if settings.deterministic_probes.on():\n                    warnings.warn(\n                        \"Deterministic probes will currently work only if you aren't training multiple independent\"\n                        \" models simultaneously.\",\n                        UserWarning,\n                    )\n                    if settings.deterministic_probes.probe_vectors is None:\n                        probe_vectors = torch.empty(matrix_shape[-1], num_random_probes, dtype=dtype, device=device)\n                        probe_vectors.bernoulli_().mul_(2).add_(-1)\n                        settings.deterministic_probes.probe_vectors = probe_vectors\n                    else:\n                        probe_vectors = settings.deterministic_probes.probe_vectors\n                else:\n                    probe_vectors = torch.empty(matrix_shape[-1], num_random_probes, dtype=dtype, device=device)\n                    probe_vectors.bernoulli_().mul_(2).add_(-1)\n\n                probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)\n                if batch_shape is not None:\n                    probe_vectors = probe_vectors.expand(*batch_shape, matrix_shape[-1], num_random_probes)\n                    probe_vector_norms = probe_vector_norms.expand(*batch_shape, 1, num_random_probes)\n            else:  # When preconditioning, probe vectors must be drawn from N(0, P)\n                if precond_lt.size()[-2:] == torch.Size([1, 1]):\n                    covar_root = precond_lt.evaluate().sqrt()\n                else:\n                    covar_root = precond_lt.root_decomposition().root\n\n                if settings.deterministic_probes.on():\n                    warnings.warn(\n                        \"Deterministic probes will currently work only if you aren't training multiple independent\"\n                        \" models simultaneously.\",\n                        UserWarning,\n                    )\n                    base_samples = settings.deterministic_probes.probe_vectors\n                    if base_samples is None or covar_root.size(-1) != base_samples.size(-2):\n                        base_samples = torch.randn(\n                            *precond_lt.batch_shape,\n                            covar_root.size(-1),\n                            num_random_probes,\n                            dtype=precond_lt.dtype,\n                            device=precond_lt.device,\n                        )\n                        settings.deterministic_probes.probe_vectors = base_samples\n\n                    probe_vectors = covar_root.matmul(base_samples).permute(-1, *range(precond_lt.dim() - 1))\n                else:\n                    base_samples = torch.randn(\n                        *precond_lt.batch_shape,\n                        covar_root.size(-1),\n                        num_random_probes,\n                        dtype=precond_lt.dtype,\n                        device=precond_lt.device,\n                    )\n                    probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)\n                probe_vectors = probe_vectors.unsqueeze(-2).transpose(0, -2).squeeze(0).transpose(-2, -1).contiguous()\n                probe_vector_norms = torch.norm(probe_vectors, p=2, dim=-2, keepdim=True)\n            probe_vectors = probe_vectors.div(probe_vector_norms)\n\n        ctx.probe_vectors = probe_vectors\n        ctx.probe_vector_norms = probe_vector_norms\n\n        if ctx.logdet and not ctx.probe_vectors.numel():\n            raise RuntimeError(\"Probe vectors were not supplied for logdet computation\")\n\n        # Collect terms for LinearCG\n        # We use LinearCG for both matrix solves and for stochastically estimating the log det\n        rhs_list = []\n        num_random_probes = 0\n        num_inv_quad_solves = 0\n\n        # RHS for logdet\n        if ctx.logdet:\n            rhs_list.append(ctx.probe_vectors)\n            num_random_probes = ctx.probe_vectors.size(-1)\n\n        # RHS for inv_quad\n        ctx.is_vector = False\n        if ctx.inv_quad:\n            if inv_quad_rhs.ndimension() == 1:\n                inv_quad_rhs = inv_quad_rhs.unsqueeze(-1)\n                ctx.is_vector = True\n            rhs_list.append(inv_quad_rhs)\n            num_inv_quad_solves = inv_quad_rhs.size(-1)\n\n        # Perform solves (for inv_quad) and tridiagonalization (for estimating logdet)\n        rhs = torch.cat(rhs_list, -1)\n        t_mat = None\n        if ctx.logdet and settings.skip_logdet_forward.off():\n            solves, t_mat = lazy_tsr._solve(rhs, preconditioner, num_tridiag=num_random_probes)\n\n        else:\n            solves = lazy_tsr._solve(rhs, preconditioner, num_tridiag=0)\n\n        # Final values to return\n        logdet_term = torch.zeros(lazy_tsr.batch_shape, dtype=ctx.dtype, device=ctx.device)\n        inv_quad_term = torch.zeros(lazy_tsr.batch_shape, dtype=ctx.dtype, device=ctx.device)\n\n        # Compute logdet from tridiagonalization\n        if ctx.logdet and settings.skip_logdet_forward.off():\n            if torch.any(torch.isnan(t_mat)).item():\n                logdet_term = torch.tensor(float(\"nan\"), dtype=ctx.dtype, device=ctx.device)\n            else:\n                if ctx.batch_shape is None:\n                    t_mat = t_mat.unsqueeze(1)\n                eigenvalues, eigenvectors = lanczos_tridiag_to_diag(t_mat)\n                slq = StochasticLQ()\n                (logdet_term,) = slq.evaluate(ctx.matrix_shape, eigenvalues, eigenvectors, [lambda x: x.log()])\n\n                # Add correction\n                if logdet_correction is not None:\n                    logdet_term = logdet_term + logdet_correction\n\n        # Extract inv_quad solves from all the solves\n        if ctx.inv_quad:\n            inv_quad_solves = solves.narrow(-1, num_random_probes, num_inv_quad_solves)\n            inv_quad_term = (inv_quad_solves * inv_quad_rhs).sum(-2)\n\n        ctx.num_random_probes = num_random_probes\n        ctx.num_inv_quad_solves = num_inv_quad_solves\n\n        to_save = list(matrix_args) + [solves]\n        ctx.save_for_backward(*to_save)\n\n        if settings.memory_efficient.off():\n            ctx._lazy_tsr = lazy_tsr\n\n        return inv_quad_term, logdet_term\n\n    @staticmethod\n    def backward(ctx, inv_quad_grad_output, logdet_grad_output):\n        matrix_arg_grads = None\n        inv_quad_rhs_grad = None\n\n        # Which backward passes should we compute?\n        compute_inv_quad_grad = inv_quad_grad_output.abs().sum() and ctx.inv_quad\n        compute_logdet_grad = logdet_grad_output.abs().sum() and ctx.logdet\n\n        # Get input arguments, and get gradients in the proper form\n        matrix_args = ctx.saved_tensors[:-1]\n        solves = ctx.saved_tensors[-1]\n\n        if hasattr(ctx, \"_lazy_tsr\"):\n            lazy_tsr = ctx._lazy_tsr\n        else:\n            lazy_tsr = ctx.representation_tree(*matrix_args)\n\n        # Fix grad_output sizes\n        if ctx.inv_quad:\n            inv_quad_grad_output = inv_quad_grad_output.unsqueeze(-2)\n        if compute_logdet_grad:\n            logdet_grad_output = logdet_grad_output.unsqueeze(-1)\n            logdet_grad_output.unsqueeze_(-1)\n\n        # Divide up the solves\n        probe_vector_solves = None\n        inv_quad_solves = None\n        neg_inv_quad_solves_times_grad_out = None\n        if compute_logdet_grad:\n            coef = 1.0 / ctx.probe_vectors.size(-1)\n            probe_vector_solves = solves.narrow(-1, 0, ctx.num_random_probes).mul(coef)\n            probe_vector_solves.mul_(ctx.probe_vector_norms).mul_(logdet_grad_output)\n            probe_vectors = ctx.probe_vectors.mul(ctx.probe_vector_norms)\n        if ctx.inv_quad:\n            inv_quad_solves = solves.narrow(-1, ctx.num_random_probes, ctx.num_inv_quad_solves)\n            neg_inv_quad_solves_times_grad_out = inv_quad_solves.mul(inv_quad_grad_output).mul_(-1)\n\n        # input_1 gradient\n        if any(ctx.needs_input_grad):\n            # Collect terms for arg grads\n            left_factors_list = []\n            right_factors_list = []\n\n            if compute_logdet_grad:\n                left_factors_list.append(probe_vector_solves)\n                if ctx.preconditioner is not None:\n                    probe_vectors = ctx.preconditioner(probe_vectors)\n                right_factors_list.append(probe_vectors)\n\n            if compute_inv_quad_grad:\n                left_factors_list.append(neg_inv_quad_solves_times_grad_out)\n                right_factors_list.append(inv_quad_solves)\n\n            left_factors = torch.cat(left_factors_list, -1)\n            right_factors = torch.cat(right_factors_list, -1)\n            matrix_arg_grads = lazy_tsr._quad_form_derivative(left_factors, right_factors)\n\n        # input_2 gradients\n        if compute_inv_quad_grad and ctx.needs_input_grad[9]:\n            inv_quad_rhs_grad = neg_inv_quad_solves_times_grad_out.mul_(-2)\n        elif ctx.inv_quad:\n            inv_quad_rhs_grad = torch.zeros_like(inv_quad_solves)\n        if ctx.is_vector:\n            inv_quad_rhs_grad.squeeze_(-1)\n\n        if ctx.inv_quad:\n            res = [inv_quad_rhs_grad] + list(matrix_arg_grads)\n        else:\n            res = list(matrix_arg_grads)\n\n        return tuple([None] * 9 + res)\n", "path": "gpytorch/functions/_inv_quad_log_det.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\n\nimport torch\nfrom torch.autograd import Function\n\nfrom .. import settings\nfrom ..utils.lanczos import lanczos_tridiag_to_diag\nfrom ..utils.stochastic_lq import StochasticLQ\n\n\nclass InvQuadLogDet(Function):\n    \"\"\"\n    Given a PSD matrix A (or a batch of PSD matrices A), this function computes one or both\n    of the following\n    - The matrix solves A^{-1} b\n    - logdet(A)\n    \"\"\"\n\n    @staticmethod\n    def forward(\n        ctx,\n        representation_tree,\n        dtype,\n        device,\n        matrix_shape,\n        batch_shape=torch.Size(),\n        inv_quad=False,\n        logdet=False,\n        probe_vectors=None,\n        probe_vector_norms=None,\n        *args,\n    ):\n        \"\"\"\n        *args - The arguments representing the PSD matrix A (or batch of PSD matrices A)\n        If self.inv_quad is true, the first entry in *args is inv_quad_rhs (Tensor)\n        - the RHS of the matrix solves.\n\n        Returns:\n        - (Scalar) The inverse quadratic form (or None, if self.inv_quad is False)\n        - (Scalar) The log determinant (or None, self.if logdet is False)\n        \"\"\"\n\n        if not (inv_quad or logdet):\n            raise RuntimeError(\"Either inv_quad or logdet must be true (or both)\")\n\n        ctx.representation_tree = representation_tree\n        ctx.dtype = dtype\n        ctx.device = device\n        ctx.matrix_shape = matrix_shape\n        ctx.batch_shape = batch_shape\n        ctx.inv_quad = inv_quad\n        ctx.logdet = logdet\n\n        matrix_args = None\n        inv_quad_rhs = None\n        if ctx.inv_quad:\n            matrix_args = args[1:]\n            inv_quad_rhs = args[0]\n        else:\n            matrix_args = args\n\n        # Get closure for matmul\n        lazy_tsr = ctx.representation_tree(*matrix_args)\n        with torch.no_grad():\n            preconditioner, precond_lt, logdet_correction = lazy_tsr._preconditioner()\n\n        ctx.preconditioner = preconditioner\n\n        if (probe_vectors is None or probe_vector_norms is None) and logdet:\n            num_random_probes = settings.num_trace_samples.value()\n            if preconditioner is None:\n                if settings.deterministic_probes.on():\n                    warnings.warn(\n                        \"Deterministic probes will currently work only if you aren't training multiple independent\"\n                        \" models simultaneously.\",\n                        UserWarning,\n                    )\n                    if settings.deterministic_probes.probe_vectors is None:\n                        probe_vectors = torch.empty(matrix_shape[-1], num_random_probes, dtype=dtype, device=device)\n                        probe_vectors.bernoulli_().mul_(2).add_(-1)\n                        settings.deterministic_probes.probe_vectors = probe_vectors\n                    else:\n                        probe_vectors = settings.deterministic_probes.probe_vectors\n                else:\n                    probe_vectors = torch.empty(matrix_shape[-1], num_random_probes, dtype=dtype, device=device)\n                    probe_vectors.bernoulli_().mul_(2).add_(-1)\n\n                probe_vector_norms = torch.norm(probe_vectors, 2, dim=-2, keepdim=True)\n                if batch_shape is not None:\n                    probe_vectors = probe_vectors.expand(*batch_shape, matrix_shape[-1], num_random_probes)\n                    probe_vector_norms = probe_vector_norms.expand(*batch_shape, 1, num_random_probes)\n            else:  # When preconditioning, probe vectors must be drawn from N(0, P)\n                if settings.deterministic_probes.on():\n                    # NOTE: calling precond_lt.root_decomposition() is expensive\n                    # because it requires Lanczos\n                    # We don't have any other choice for when we want to use deterministic probes, however\n                    if precond_lt.size()[-2:] == torch.Size([1, 1]):\n                        covar_root = precond_lt.evaluate().sqrt()\n                    else:\n                        covar_root = precond_lt.root_decomposition().root\n\n                    warnings.warn(\n                        \"Deterministic probes will currently work only if you aren't training multiple independent\"\n                        \" models simultaneously.\",\n                        UserWarning,\n                    )\n                    base_samples = settings.deterministic_probes.probe_vectors\n                    if base_samples is None or covar_root.size(-1) != base_samples.size(-2):\n                        base_samples = torch.randn(\n                            *precond_lt.batch_shape,\n                            covar_root.size(-1),\n                            num_random_probes,\n                            dtype=precond_lt.dtype,\n                            device=precond_lt.device,\n                        )\n                        settings.deterministic_probes.probe_vectors = base_samples\n\n                    probe_vectors = covar_root.matmul(base_samples).permute(-1, *range(precond_lt.dim() - 1))\n                else:\n                    probe_vectors = precond_lt.zero_mean_mvn_samples(num_random_probes)\n                probe_vectors = probe_vectors.unsqueeze(-2).transpose(0, -2).squeeze(0).transpose(-2, -1).contiguous()\n                probe_vector_norms = torch.norm(probe_vectors, p=2, dim=-2, keepdim=True)\n            probe_vectors = probe_vectors.div(probe_vector_norms)\n\n        ctx.probe_vectors = probe_vectors\n        ctx.probe_vector_norms = probe_vector_norms\n\n        if ctx.logdet and not ctx.probe_vectors.numel():\n            raise RuntimeError(\"Probe vectors were not supplied for logdet computation\")\n\n        # Collect terms for LinearCG\n        # We use LinearCG for both matrix solves and for stochastically estimating the log det\n        rhs_list = []\n        num_random_probes = 0\n        num_inv_quad_solves = 0\n\n        # RHS for logdet\n        if ctx.logdet:\n            rhs_list.append(ctx.probe_vectors)\n            num_random_probes = ctx.probe_vectors.size(-1)\n\n        # RHS for inv_quad\n        ctx.is_vector = False\n        if ctx.inv_quad:\n            if inv_quad_rhs.ndimension() == 1:\n                inv_quad_rhs = inv_quad_rhs.unsqueeze(-1)\n                ctx.is_vector = True\n            rhs_list.append(inv_quad_rhs)\n            num_inv_quad_solves = inv_quad_rhs.size(-1)\n\n        # Perform solves (for inv_quad) and tridiagonalization (for estimating logdet)\n        rhs = torch.cat(rhs_list, -1)\n        t_mat = None\n        if ctx.logdet and settings.skip_logdet_forward.off():\n            solves, t_mat = lazy_tsr._solve(rhs, preconditioner, num_tridiag=num_random_probes)\n\n        else:\n            solves = lazy_tsr._solve(rhs, preconditioner, num_tridiag=0)\n\n        # Final values to return\n        logdet_term = torch.zeros(lazy_tsr.batch_shape, dtype=ctx.dtype, device=ctx.device)\n        inv_quad_term = torch.zeros(lazy_tsr.batch_shape, dtype=ctx.dtype, device=ctx.device)\n\n        # Compute logdet from tridiagonalization\n        if ctx.logdet and settings.skip_logdet_forward.off():\n            if torch.any(torch.isnan(t_mat)).item():\n                logdet_term = torch.tensor(float(\"nan\"), dtype=ctx.dtype, device=ctx.device)\n            else:\n                if ctx.batch_shape is None:\n                    t_mat = t_mat.unsqueeze(1)\n                eigenvalues, eigenvectors = lanczos_tridiag_to_diag(t_mat)\n                slq = StochasticLQ()\n                (logdet_term,) = slq.evaluate(ctx.matrix_shape, eigenvalues, eigenvectors, [lambda x: x.log()])\n\n                # Add correction\n                if logdet_correction is not None:\n                    logdet_term = logdet_term + logdet_correction\n\n        # Extract inv_quad solves from all the solves\n        if ctx.inv_quad:\n            inv_quad_solves = solves.narrow(-1, num_random_probes, num_inv_quad_solves)\n            inv_quad_term = (inv_quad_solves * inv_quad_rhs).sum(-2)\n\n        ctx.num_random_probes = num_random_probes\n        ctx.num_inv_quad_solves = num_inv_quad_solves\n\n        to_save = list(matrix_args) + [solves]\n        ctx.save_for_backward(*to_save)\n\n        if settings.memory_efficient.off():\n            ctx._lazy_tsr = lazy_tsr\n\n        return inv_quad_term, logdet_term\n\n    @staticmethod\n    def backward(ctx, inv_quad_grad_output, logdet_grad_output):\n        matrix_arg_grads = None\n        inv_quad_rhs_grad = None\n\n        # Which backward passes should we compute?\n        compute_inv_quad_grad = inv_quad_grad_output.abs().sum() and ctx.inv_quad\n        compute_logdet_grad = logdet_grad_output.abs().sum() and ctx.logdet\n\n        # Get input arguments, and get gradients in the proper form\n        matrix_args = ctx.saved_tensors[:-1]\n        solves = ctx.saved_tensors[-1]\n\n        if hasattr(ctx, \"_lazy_tsr\"):\n            lazy_tsr = ctx._lazy_tsr\n        else:\n            lazy_tsr = ctx.representation_tree(*matrix_args)\n\n        # Fix grad_output sizes\n        if ctx.inv_quad:\n            inv_quad_grad_output = inv_quad_grad_output.unsqueeze(-2)\n        if compute_logdet_grad:\n            logdet_grad_output = logdet_grad_output.unsqueeze(-1)\n            logdet_grad_output.unsqueeze_(-1)\n\n        # Divide up the solves\n        probe_vector_solves = None\n        inv_quad_solves = None\n        neg_inv_quad_solves_times_grad_out = None\n        if compute_logdet_grad:\n            coef = 1.0 / ctx.probe_vectors.size(-1)\n            probe_vector_solves = solves.narrow(-1, 0, ctx.num_random_probes).mul(coef)\n            probe_vector_solves.mul_(ctx.probe_vector_norms).mul_(logdet_grad_output)\n            probe_vectors = ctx.probe_vectors.mul(ctx.probe_vector_norms)\n        if ctx.inv_quad:\n            inv_quad_solves = solves.narrow(-1, ctx.num_random_probes, ctx.num_inv_quad_solves)\n            neg_inv_quad_solves_times_grad_out = inv_quad_solves.mul(inv_quad_grad_output).mul_(-1)\n\n        # input_1 gradient\n        if any(ctx.needs_input_grad):\n            # Collect terms for arg grads\n            left_factors_list = []\n            right_factors_list = []\n\n            if compute_logdet_grad:\n                left_factors_list.append(probe_vector_solves)\n                if ctx.preconditioner is not None:\n                    probe_vectors = ctx.preconditioner(probe_vectors)\n                right_factors_list.append(probe_vectors)\n\n            if compute_inv_quad_grad:\n                left_factors_list.append(neg_inv_quad_solves_times_grad_out)\n                right_factors_list.append(inv_quad_solves)\n\n            left_factors = torch.cat(left_factors_list, -1)\n            right_factors = torch.cat(right_factors_list, -1)\n            matrix_arg_grads = lazy_tsr._quad_form_derivative(left_factors, right_factors)\n\n        # input_2 gradients\n        if compute_inv_quad_grad and ctx.needs_input_grad[9]:\n            inv_quad_rhs_grad = neg_inv_quad_solves_times_grad_out.mul_(-2)\n        elif ctx.inv_quad:\n            inv_quad_rhs_grad = torch.zeros_like(inv_quad_solves)\n        if ctx.is_vector:\n            inv_quad_rhs_grad.squeeze_(-1)\n\n        if ctx.inv_quad:\n            res = [inv_quad_rhs_grad] + list(matrix_arg_grads)\n        else:\n            res = list(matrix_arg_grads)\n\n        return tuple([None] * 9 + res)\n", "path": "gpytorch/functions/_inv_quad_log_det.py"}]}