# Buggy code
def _preconditioner(self):
    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

    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
        )
        if torch.any(torch.isnan(self._piv_chol_self)).item():
            warnings.warn(
                "NaNs encountered in preconditioner computation. Attempting to continue without preconditioning.",
                NumericalWarning,
            )
            return None, None, None
        self._init_cache()

    # NOTE: We cannot memoize this precondition closure as it causes a memory leak
    def precondition_closure(tensor):
        qqt = self._q_cache.matmul(self._q_cache.transpose(-2, -1).matmul(tensor))
        if self._constant_diag:
            return (1 / self._noise) * (tensor - qqt)
        return (tensor / self._noise) - qqt

    return (precondition_closure, self._precond_lt, self._precond_logdet_cache)


# Traceback (RuntimeError)
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

# Fixed code
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
        )
        if torch.any(torch.isnan(self._piv_chol_self)).item():
            warnings.warn(
                "NaNs encountered in preconditioner computation. Attempting to continue without preconditioning.",
                NumericalWarning,
            )
            return None, None, None
        self._init_cache()

    # 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)
        return (tensor / self._noise) - qqt

    return (precondition_closure, self._precond_lt, self._precond_logdet_cache)