Repository: cornellius-gp/gpytorch (python)

Issue:
[Bug] NUTS not working when GammaPrior in model
# 🐛 Bug

I'm trying to use the [NUTS example](https://gpytorch.readthedocs.io/en/latest/examples/01_Simple_GP_Regression/Simple_GP_Regression_Fully_Bayesian.html#Sampling-Hyperparamters-with-GPyTorch-+-NUTS) together with the model from [this example](https://gpytorch.readthedocs.io/en/latest/examples/00_Basic_Usage/Hyperparameters.html#Putting-it-Together) and get 
```
AttributeError: 'GammaPrior' object has no attribute 'concentration'
```

## To reproduce

** Code snippet to reproduce **
```python
class ExactGPWithPriors(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPWithPriors, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()

        lengthscale_prior = gpytorch.priors.GammaPrior(3.0, 6.0)
        outputscale_prior = gpytorch.priors.GammaPrior(2.0, 0.15)
            
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(
                lengthscale_prior=lengthscale_prior,
            ),
            outputscale_prior=outputscale_prior
        )

        # Initialize lengthscale and outputscale to mean of priors
        self.covar_module.base_kernel.lengthscale = lengthscale_prior.mean
        self.covar_module.outputscale = outputscale_prior.mean

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood(
    noise_constraint=gpytorch.constraints.GreaterThan(1e-2),
)

model = ExactGPWithPriors(train_x, train_y, likelihood)

mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

def pyro_model(x, y):
    model.pyro_sample_from_prior()
    output = model(x)
    loss = mll.pyro_factor(output, y)
    return y

nuts_kernel = pyro.infer.mcmc.NUTS(pyro_model, adapt_step_size=True)
mcmc_run = pyro.infer.mcmc.MCMC(nuts_kernel, num_samples=100, warmup_steps=200)
mcmc_run.run(train_x, train_y)
```

** Stack trace/error message **
```
Warmup:   0%|          | 0/300 [00:00, ?it/s]
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-22-1dcdc695642b> in <module>
      9 nuts_kernel = pyro.infer.mcmc.NUTS(pyro_model, adapt_step_size=True)
     10 mcmc_run = pyro.infer.mcmc.MCMC(nuts_kernel, num_samples=100, warmup_steps=200)
---> 11 mcmc_run.run(train_x, train_y)

~/.local/lib/python3.8/site-packages/pyro/poutine/messenger.py in _context_wrap(context, fn, *args, **kwargs)
      6 def _context_wrap(context, fn, *args, **kwargs):
      7     with context:
----> 8         return fn(*args, **kwargs)
      9 
     10 

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/api.py in run(self, *args, **kwargs)
    352         z_flat_acc = [[] for _ in range(self.num_chains)]
    353         with pyro.validation_enabled(not self.disable_validation):
--> 354             for x, chain_id in self.sampler.run(*args, **kwargs):
    355                 if num_samples[chain_id] == 0:
    356                     num_samples[chain_id] += 1

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/api.py in run(self, *args, **kwargs)
    161             logger = initialize_logger(logger, "", progress_bar)
    162             hook_w_logging = _add_logging_hook(logger, progress_bar, self.hook)
--> 163             for sample in _gen_samples(self.kernel, self.warmup_steps, self.num_samples, hook_w_logging,
    164                                        i if self.num_chains > 1 else None,
    165                                        *args, **kwargs):

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/api.py in _gen_samples(kernel, warmup_steps, num_samples, hook, chain_id, *args, **kwargs)
    105 
    106 def _gen_samples(kernel, warmup_steps, num_samples, hook, chain_id, *args, **kwargs):
--> 107     kernel.setup(warmup_steps, *args, **kwargs)
    108     params = kernel.initial_params
    109     # yield structure (key, value.shape) of params

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/hmc.py in setup(self, warmup_steps, *args, **kwargs)
    258         self._warmup_steps = warmup_steps
    259         if self.model is not None:
--> 260             self._initialize_model_properties(args, kwargs)
    261         potential_energy = self.potential_fn(self.initial_params)
    262         self._cache(self.initial_params, potential_energy, None)

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/hmc.py in _initialize_model_properties(self, model_args, model_kwargs)
    223 
    224     def _initialize_model_properties(self, model_args, model_kwargs):
--> 225         init_params, potential_fn, transforms, trace = initialize_model(
    226             self.model,
    227             model_args,

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/util.py in initialize_model(model, model_args, model_kwargs, transforms, max_plate_nesting, jit_compile, jit_options, skip_jit_warnings, num_chains)
    363         automatic_transform_enabled = False
    364     if max_plate_nesting is None:
--> 365         max_plate_nesting = _guess_max_plate_nesting(model, model_args, model_kwargs)
    366     # Wrap model in `poutine.enum` to enumerate over discrete latent sites.
    367     # No-op if model does not have any discrete latents.

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/util.py in _guess_max_plate_nesting(model, args, kwargs)
    231     """
    232     with poutine.block():
--> 233         model_trace = poutine.trace(model).get_trace(*args, **kwargs)
    234     sites = [site for site in model_trace.nodes.values()
    235              if site["type"] == "sample"]

~/.local/lib/python3.8/site-packages/pyro/poutine/trace_messenger.py in get_trace(self, *args, **kwargs)
    161         Calls this poutine and returns its trace instead of the function's return value.
    162         """
--> 163         self(*args, **kwargs)
    164         return self.msngr.get_trace()

~/.local/lib/python3.8/site-packages/pyro/poutine/trace_messenger.py in __call__(self, *args, **kwargs)
    141                                       args=args, kwargs=kwargs)
    142             try:
--> 143                 ret = self.fn(*args, **kwargs)
    144             except (ValueError, RuntimeError):
    145                 exc_type, exc_value, traceback = sys.exc_info()

<ipython-input-22-1dcdc695642b> in pyro_model(x, y)
      2 
      3 def pyro_model(x, y):
----> 4     model.pyro_sample_from_prior()
      5     output = model(x)
      6     loss = mll.pyro_factor(output, y)

~/.local/lib/python3.8/site-packages/gpytorch/module.py in pyro_sample_from_prior(self)
    287         parameters of the model that have GPyTorch priors registered to them.
    288         """
--> 289         return _pyro_sample_from_prior(module=self, memo=None, prefix="")
    290 
    291     def local_load_samples(self, samples_dict, memo, prefix):

~/.local/lib/python3.8/site-packages/gpytorch/module.py in _pyro_sample_from_prior(module, memo, prefix)
    396     for mname, module_ in module.named_children():
    397         submodule_prefix = prefix + ("." if prefix else "") + mname
--> 398         _pyro_sample_from_prior(module=module_, memo=memo, prefix=submodule_prefix)
    399 
    400 

~/.local/lib/python3.8/site-packages/gpytorch/module.py in _pyro_sample_from_prior(module, memo, prefix)
    391                 memo.add(prior)
    392                 prior = prior.expand(closure().shape)
--> 393                 value = pyro.sample(prefix + ("." if prefix else "") + prior_name, prior)
    394                 setting_closure(value)
    395 

~/.local/lib/python3.8/site-packages/pyro/primitives.py in sample(name, fn, *args, **kwargs)
    108             msg["is_observed"] = True
    109         # apply the stack and return its return value
--> 110         apply_stack(msg)
    111         return msg["value"]
    112 

~/.local/lib/python3.8/site-packages/pyro/poutine/runtime.py in apply_stack(initial_msg)
    193             break
    194 
--> 195     default_process_message(msg)
    196 
    197     for frame in stack[-pointer:]:

~/.local/lib/python3.8/site-packages/pyro/poutine/runtime.py in default_process_message(msg)
    154         return msg
    155 
--> 156     msg["value"] = msg["fn"](*msg["args"], **msg["kwargs"])
    157 
    158     # after fn has been called, update msg to prevent it from being called again.

~/.local/lib/python3.8/site-packages/pyro/distributions/torch_distribution.py in __call__(self, sample_shape)
     38         :rtype: torch.Tensor
     39         """
---> 40         return self.rsample(sample_shape) if self.has_rsample else self.sample(sample_shape)
     41 
     42     @property

/usr/lib/python3.8/site-packages/torch/distributions/gamma.py in rsample(self, sample_shape)
     59     def rsample(self, sample_shape=torch.Size()):
     60         shape = self._extended_shape(sample_shape)
---> 61         value = _standard_gamma(self.concentration.expand(shape)) / self.rate.expand(shape)
     62         value.detach().clamp_(min=torch.finfo(value.dtype).tiny)  # do not record in autograd graph
     63         return value

/usr/lib/python3.8/site-packages/torch/nn/modules/module.py in __getattr__(self, name)
    582             if name in modules:
    583                 return modules[name]
--> 584         raise AttributeError("'{}' object has no attribute '{}'".format(
    585             type(self).__name__, name))
    586 

AttributeError: 'GammaPrior' object has no attribute 'concentration'
```

## System information
gpytorch 0.3.6 (pip install git, https://github.com/cornellius-gp/gpytorch/commit/4a1ba02d2367e4e9dd03eb1ccbfa4707da02dd08)
torch 1.3.1
pyro 1.0.0

Pull request description:
Custom expand methods for priors
Inheriting the behavior of `expand` from the base torch distributions was causing really strange behavior, I think because our priors also inherit from Module. In addition to the bug in #976, calling `prior.expand(some_torch_size)` was actually modifying `prior` in place in addition to returning a new Module with the correct size.

These issues were easily fixed by just defining our own `expand` methods for the various priors. I also think that this is a rare setting where overriding and defining our own expand behavior will actually make things easier to debug, because the upstream behavior is extremely weird: https://github.com/pytorch/pytorch/blob/dd52f50fc85e6710f020936c1fc5f14673508350/torch/distributions/gamma.py#L50-L57

Fixes #976 

Patch:
diff --git a/.pylintrc b/.pylintrc
deleted file mode 100644
index 317e745c..00000000
--- a/.pylintrc
+++ /dev/null
@@ -1,5 +0,0 @@
-[TYPECHECK]
-
-# List of members which are set dynamically and missed by Pylint inference
-# system, and so shouldn't trigger E1101 when accessed.
-generated-members=numpy.*, torch.*
diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py
index fe25d4ea..9fa5b451 100644
--- a/gpytorch/lazy/added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/added_diag_lazy_tensor.py
@@ -5,7 +5,7 @@ import warnings
 import torch
 
 from .. import settings
-from ..utils import broadcasting, pivoted_cholesky
+from ..utils import broadcasting, pivoted_cholesky, woodbury
 from .diag_lazy_tensor import DiagLazyTensor
 from .psd_sum_lazy_tensor import PsdSumLazyTensor
 from .root_lazy_tensor import RootLazyTensor
@@ -27,7 +27,7 @@ class AddedDiagLazyTensor(SumLazyTensor):
         broadcasting._mul_broadcast_shape(lazy_tensors[0].shape, lazy_tensors[1].shape)
 
         if isinstance(lazy_tensors[0], DiagLazyTensor) and isinstance(lazy_tensors[1], DiagLazyTensor):
-            raise RuntimeError("Trying to lazily add two DiagLazyTensors. Create a single DiagLazyTensor instead.")
+            raise RuntimeError("Trying to lazily add two DiagLazyTensors. " "Create a single DiagLazyTensor instead.")
         elif isinstance(lazy_tensors[0], DiagLazyTensor):
             self._diag_tensor = lazy_tensors[0]
             self._lazy_tensor = lazy_tensors[1]
@@ -39,15 +39,6 @@ class AddedDiagLazyTensor(SumLazyTensor):
 
         self.preconditioner_override = preconditioner_override
 
-        # Placeholders
-        self._constant_diag = None
-        self._noise = None
-        self._piv_chol_self = None  # <- Doesn't need to be an attribute, but used for testing purposes
-        self._precond_lt = None
-        self._precond_logdet_cache = None
-        self._q_cache = None
-        self._r_cache = None
-
     def _matmul(self, rhs):
         return torch.addcmul(self._lazy_tensor._matmul(rhs), self._diag_tensor._diag.unsqueeze(-1), rhs)
 
@@ -69,7 +60,7 @@ class AddedDiagLazyTensor(SumLazyTensor):
         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:
+        if not (hasattr(self, "_woodbury_cache") or hasattr(self, "self._q_cache")):
             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():
@@ -77,46 +68,66 @@ class AddedDiagLazyTensor(SumLazyTensor):
                     "NaNs encountered in preconditioner computation. Attempting to continue without preconditioning."
                 )
                 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.t().matmul(tensor))
-            if self._constant_diag:
-                return (1 / self._noise) * (tensor - qqt)
-            return (tensor / self._noise) - qqt
+            if self._piv_chol_self.dim() == 2:  # TODO: Whenever PyTorch supports batch mode
+                *batch_shape, n, k = self._piv_chol_self.shape
+                self._noise = self._diag_tensor.diag().unsqueeze(-1)
 
-        return (precondition_closure, self._precond_lt, self._precond_logdet_cache)
+                self.constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))
 
-    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))
-        eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)
+                eye = torch.eye(k, dtype=self._piv_chol_self.dtype, device=self._piv_chol_self.device)
 
-        if self._constant_diag:
-            self._init_cache_for_constant_diag(eye, batch_shape, n, k)
-        else:
-            self._init_cache_for_non_constant_diag(eye, batch_shape, n)
+                if self.constant_diag:
+                    # We can factor out the noise for for both QR and solves.
+                    self.noise_constant = self._noise[0].squeeze()
+                    self._q_cache, self._r_cache = torch.qr(
+                        torch.cat((self._piv_chol_self, self.noise_constant.sqrt() * eye))
+                    )
+                    self._q_cache = self._q_cache[:n, :]
+
+                    # 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 = logdet + (n - k) * self.noise_constant.log()
+                    self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
 
-        self._precond_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)
+                else:
+                    # 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._q_cache[:n, :] / self._noise.sqrt()
 
-    def _init_cache_for_constant_diag(self, eye, batch_shape, n, k):
-        # We can factor out the noise for for both QR and solves.
-        self._noise = self._noise[0].squeeze()
-        self._q_cache, self._r_cache = torch.qr(torch.cat((self._piv_chol_self, self._noise.sqrt() * eye)))
-        self._q_cache = self._q_cache[:n, :]
+                    logdet = self._r_cache.diagonal(dim1=-1, dim2=-2).abs().log().sum(-1).mul(2) - (
+                        1.0 / self._noise
+                    ).log().sum([-1, -2])
+                    self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
 
-        # 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 = logdet + (n - k) * self._noise.log()
-        self._precond_logdet_cache = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
+            else:
+                self._woodbury_cache, self._inv_scale, self._precond_logdet_cache = woodbury.woodbury_factor(
+                    self._piv_chol_self, self._piv_chol_self, self._diag_tensor.diag(), logdet=True
+                )
+                self._scaled_inv_diag = self._inv_scale / self._diag_tensor.diag().unsqueeze(-1)
+                self._scaled_inv_diag_piv_chol_self = self._piv_chol_self * self._scaled_inv_diag
+
+            self.preconditioner_lt = PsdSumLazyTensor(RootLazyTensor(self._piv_chol_self), self._diag_tensor)
 
-    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._q_cache[:n, :] / self._noise.sqrt()
+        # NOTE to future self:
+        # We cannot memoize this precondition closure
+        # It causes a memory leak otherwise
+        def precondition_closure(tensor):
+            if hasattr(self, "_q_cache"):
+                if self.constant_diag:
+                    return (1 / self.noise_constant) * (tensor - self._q_cache.matmul(self._q_cache.t().matmul(tensor)))
+
+                else:
+                    return (tensor / self._noise) - self._q_cache.matmul(self._q_cache.t().matmul(tensor))
+
+            else:
+                res = woodbury.woodbury_solve(
+                    tensor,
+                    self._scaled_inv_diag_piv_chol_self,
+                    self._woodbury_cache,
+                    self._scaled_inv_diag,
+                    self._inv_scale,
+                )
+                return res
 
-        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()
+        return (precondition_closure, self.preconditioner_lt, self._precond_logdet_cache)
diff --git a/gpytorch/priors/horseshoe_prior.py b/gpytorch/priors/horseshoe_prior.py
index d081e428..65845ddc 100644
--- a/gpytorch/priors/horseshoe_prior.py
+++ b/gpytorch/priors/horseshoe_prior.py
@@ -7,7 +7,6 @@ import torch
 from torch.distributions import HalfCauchy, Normal, constraints
 from torch.nn import Module as TModule
 
-from gpytorch.distributions import Distribution
 from gpytorch.priors.prior import Prior
 
 
@@ -56,9 +55,5 @@ class HorseshoePrior(Prior):
         return param_sample
 
     def expand(self, expand_shape, _instance=None):
-        new = self._get_checked_instance(HorseshoePrior)
         batch_shape = torch.Size(expand_shape)
-        new.scale = self.scale.expand(batch_shape)
-        super(Distribution, new).__init__(batch_shape)
-        new._validate_args = self._validate_args
-        return new
+        return HorseshoePrior(self.scale.expand(batch_shape))
diff --git a/gpytorch/priors/torch_priors.py b/gpytorch/priors/torch_priors.py
index 97c972a4..81e1c9d0 100644
--- a/gpytorch/priors/torch_priors.py
+++ b/gpytorch/priors/torch_priors.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import torch
 from torch.distributions import Gamma, LogNormal, MultivariateNormal, Normal, Uniform
 from torch.nn import Module as TModule
 
@@ -25,7 +26,8 @@ class NormalPrior(Prior, Normal):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return Normal.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return NormalPrior(self.loc.expand(batch_shape), self.scale.expand(batch_shape))
 
 
 class LogNormalPrior(Prior, LogNormal):
@@ -39,7 +41,8 @@ class LogNormalPrior(Prior, LogNormal):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return LogNormal.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return LogNormalPrior(self.loc.expand(batch_shape), self.scale.expand(batch_shape))
 
 
 class UniformPrior(Prior, Uniform):
@@ -53,7 +56,8 @@ class UniformPrior(Prior, Uniform):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return Uniform.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return UniformPrior(self.low.expand(batch_shape), self.high.expand(batch_shape))
 
 
 class GammaPrior(Prior, Gamma):
@@ -71,7 +75,11 @@ class GammaPrior(Prior, Gamma):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return Gamma.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return GammaPrior(self.concentration.expand(batch_shape), self.rate.expand(batch_shape))
+
+    def __call__(self, *args, **kwargs):
+        return super(Gamma, self).__call__(*args, **kwargs)
 
 
 class MultivariateNormalPrior(Prior, MultivariateNormal):
diff --git a/gpytorch/utils/woodbury.py b/gpytorch/utils/woodbury.py
new file mode 100644
index 00000000..a864fc93
--- /dev/null
+++ b/gpytorch/utils/woodbury.py
@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+
+import torch
+
+from .. import settings
+
+
+def woodbury_factor(umat, vmat, diag, logdet=False):
+    r"""
+    Given a matrix factorized as :math:`(D + U V^T)`, where
+    :math:`U`, :math:`V` are (n x k) and :math:`D` is a (n x n) diagonal matrix,
+    returns the matrix :math:`R` so that
+
+    .. math::
+        \begin{equation*}
+            R = (I_k + V^T D^-1 U)^{-1} V^T
+        \end{equation*}
+
+    to be used in solves with (D + U V^T) via the Woodbury formula.
+    Can also be used in batch mode, where U, V, and D are batches of matrices
+
+    If the :attr:`logdet=True` flag is set, this function will also return the
+    log determinant of :math:`(D + U V^T)`.
+
+    Args:
+        :attr:`umat` (Tensor n x k):
+            The left matrix factor
+        :attr:`vmat` (Tensor n x k):
+            The right matrix factor
+        :attr:`diag` (Tensor n):
+            The diagonal of D
+        :attr:`logdet` (bool):
+            Whether or not to return the log determinant
+
+    Returns:
+        Tensor (k x n) (if `logdet=False`)
+        Tensor (k x n), Tensor () (if `logdet=True`)
+    """
+    if settings.debug.on():
+        if umat.shape != vmat.shape:
+            raise ValueError("umat ({}) and vmat ({}) must have the same shape.".format(umat.shape, vmat.shape))
+        if umat.shape[:-1] != diag.shape:
+            raise ValueError("Incompatible shape for diag ({}) given umat shape ({}).".format(diag.shape, umat.shape))
+
+    # Sizes
+    *batch_shape, n, k = umat.shape
+
+    # These reshapes make it easier to use faster blas calls
+    umat = umat.view(-1, n, k)
+    vmat = vmat.view(-1, n, k)
+    diag = diag.view(-1, n, 1)
+
+    # Scale the diagonal
+    # s = scale = max |1 / diag|
+    inv_scale = diag.abs().min()
+    scaled_inv_diag = inv_scale / diag
+
+    # Compute (1/s (I_k + V^T D^-1 U)), where s is a scale factor
+    inner_mat = torch.baddbmm(
+        inv_scale,
+        torch.eye(k, dtype=scaled_inv_diag.dtype, device=scaled_inv_diag.device),
+        1,
+        vmat.transpose(-1, -2),
+        umat * scaled_inv_diag,
+    )
+
+    # Compute the cholesky factor of (1/s (I_k + V^T D^-1 U))
+    chol = torch.cholesky(inner_mat)
+
+    # Compute s (I_k + V^T D^-1 U)^-1 V^T
+    R = torch.cholesky_solve(vmat.transpose(-1, -2), chol).view(*batch_shape, k, n)
+
+    # Maybe compute the log determinant
+    if logdet:
+        # Using the matrix determinant lemma here
+        logdet = chol.diagonal(dim1=-1, dim2=-2).log().sum(-1).mul(2) - scaled_inv_diag.log().sum([-1, -2])
+        # Undo the effect of scaling on D^{-1}
+        logdet = logdet + inv_scale.log().mul(n - k)
+        # Reshape
+        logdet = logdet.view(*batch_shape) if len(batch_shape) else logdet.squeeze()
+        return R, inv_scale, logdet
+    else:
+        return R, inv_scale
+
+
+def woodbury_solve(rhs, scaled_inv_diag_umat, woodbury_factor, scaled_inv_diag, inv_scale):
+    """
+    Solves the system of equations: :math:`(D + U V^T)x = b` using the Woodbury formula,
+    where x is the right-hand-side (size n), U, V are (n x k), and D is a (n x n) diagonal matrix.
+    Can also be used in batch mode, where U, V, and D are batches of matrices and rhs is a batch of right-hand-side.
+
+    This should be used after calling woodbury_factor.
+
+    Args:
+        :attr:`rhs` (size n x t)
+            Right hand side vector b to solve with.
+        :attr:`scaled_inv_diag_umat` (n x k)
+            The `D^{-1} U` matrix
+        :attr:`woodbury_factor` (n x k)
+            The result of calling woodbury_factor on U, V, and D.
+        :attr:`diag` (vector)
+            The diagonal of D
+    """
+    # (D + UV^T)x = D^-1 x - D^-1 U ((I + V^T D^-1 U)^-1 V^T) D^-1 x
+    #             = D^-1 x - D^-1 U (1/s (woodbury_factor)) D^-1 x
+    #             = s( E^-1 x - E^-1 U (woodbury_factor) E^-1 x )
+    scaled_inv_diag_rhs = rhs * scaled_inv_diag
+    res = torch.add(scaled_inv_diag_rhs, -1, scaled_inv_diag_umat @ (woodbury_factor @ scaled_inv_diag_rhs))
+    res = res.div_(inv_scale)
+
+    # Reshape the result to be the correct shape
+    return res


Test patch:
diff --git a/test/utils/test_pivoted_cholesky.py b/test/utils/test_pivoted_cholesky.py
index 524dbe37..a75839ab 100644
--- a/test/utils/test_pivoted_cholesky.py
+++ b/test/utils/test_pivoted_cholesky.py
@@ -10,7 +10,7 @@ import torch
 from gpytorch import settings
 from gpytorch.kernels import RBFKernel
 from gpytorch.test.utils import approx_equal
-from gpytorch.utils import pivoted_cholesky
+from gpytorch.utils import pivoted_cholesky, woodbury
 
 
 class TestPivotedCholesky(unittest.TestCase):
@@ -34,6 +34,24 @@ class TestPivotedCholesky(unittest.TestCase):
         covar_approx = piv_chol @ piv_chol.transpose(-1, -2)
         self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
 
+    def test_solve(self):
+        size = 100
+        train_x = torch.linspace(0, 1, size)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
+        piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
+        woodbury_factor, inv_scale, logdet = woodbury.woodbury_factor(piv_chol, piv_chol, torch.ones(100), logdet=True)
+        self.assertTrue(approx_equal(logdet, (piv_chol @ piv_chol.transpose(-1, -2) + torch.eye(100)).logdet(), 2e-4))
+
+        rhs_vector = torch.randn(100, 50)
+        shifted_covar_matrix = covar_matrix + torch.eye(size)
+        real_solve = shifted_covar_matrix.inverse().matmul(rhs_vector)
+        scaled_inv_diag = (inv_scale / torch.ones(100)).unsqueeze(-1)
+        approx_solve = woodbury.woodbury_solve(
+            rhs_vector, piv_chol * scaled_inv_diag, woodbury_factor, scaled_inv_diag, inv_scale
+        )
+
+        self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
+
     def test_solve_qr(self, dtype=torch.float64, tol=1e-8):
         size = 50
         X = torch.rand((size, 2)).to(dtype=dtype)
@@ -106,6 +124,37 @@ class TestPivotedCholeskyBatch(unittest.TestCase):
 
         self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
 
+    def test_solve(self):
+        size = 100
+        train_x = torch.cat(
+            [torch.linspace(0, 1, size).unsqueeze(0), torch.linspace(0, 0.5, size).unsqueeze(0)], 0
+        ).unsqueeze(-1)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
+        piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
+        woodbury_factor, inv_scale, logdet = woodbury.woodbury_factor(
+            piv_chol, piv_chol, torch.ones(2, 100), logdet=True
+        )
+        actual_logdet = torch.stack(
+            [mat.logdet() for mat in (piv_chol @ piv_chol.transpose(-1, -2) + torch.eye(100)).view(-1, 100, 100)], 0
+        ).view(2)
+        self.assertTrue(approx_equal(logdet, actual_logdet, 2e-4))
+
+        rhs_vector = torch.randn(2, 100, 5)
+        shifted_covar_matrix = covar_matrix + torch.eye(size)
+        real_solve = torch.cat(
+            [
+                shifted_covar_matrix[0].inverse().matmul(rhs_vector[0]).unsqueeze(0),
+                shifted_covar_matrix[1].inverse().matmul(rhs_vector[1]).unsqueeze(0),
+            ],
+            0,
+        )
+        scaled_inv_diag = (inv_scale / torch.ones(2, 100)).unsqueeze(-1)
+        approx_solve = woodbury.woodbury_solve(
+            rhs_vector, piv_chol * scaled_inv_diag, woodbury_factor, scaled_inv_diag, inv_scale
+        )
+
+        self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
+
 
 class TestPivotedCholeskyMultiBatch(unittest.TestCase):
     def setUp(self):
@@ -145,6 +194,61 @@ class TestPivotedCholeskyMultiBatch(unittest.TestCase):
 
         self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
 
+    def test_solve(self):
+        size = 100
+        train_x = torch.cat(
+            [
+                torch.linspace(0, 1, size).unsqueeze(0),
+                torch.linspace(0, 0.5, size).unsqueeze(0),
+                torch.linspace(0, 0.25, size).unsqueeze(0),
+                torch.linspace(0, 1.25, size).unsqueeze(0),
+                torch.linspace(0, 1.5, size).unsqueeze(0),
+                torch.linspace(0, 1, size).unsqueeze(0),
+                torch.linspace(0, 0.5, size).unsqueeze(0),
+                torch.linspace(0, 0.25, size).unsqueeze(0),
+                torch.linspace(0, 1.25, size).unsqueeze(0),
+                torch.linspace(0, 1.25, size).unsqueeze(0),
+                torch.linspace(0, 1.5, size).unsqueeze(0),
+                torch.linspace(0, 1, size).unsqueeze(0),
+            ],
+            0,
+        ).unsqueeze(-1)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate().view(2, 2, 3, size, size)
+        piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
+        woodbury_factor, inv_scale, logdet = woodbury.woodbury_factor(
+            piv_chol, piv_chol, torch.ones(2, 2, 3, 100), logdet=True
+        )
+        actual_logdet = torch.stack(
+            [mat.logdet() for mat in (piv_chol @ piv_chol.transpose(-1, -2) + torch.eye(100)).view(-1, 100, 100)], 0
+        ).view(2, 2, 3)
+        self.assertTrue(approx_equal(logdet, actual_logdet, 2e-4))
+
+        rhs_vector = torch.randn(2, 2, 3, 100, 5)
+        shifted_covar_matrix = covar_matrix + torch.eye(size)
+        real_solve = torch.cat(
+            [
+                shifted_covar_matrix[0, 0, 0].inverse().matmul(rhs_vector[0, 0, 0]).unsqueeze(0),
+                shifted_covar_matrix[0, 0, 1].inverse().matmul(rhs_vector[0, 0, 1]).unsqueeze(0),
+                shifted_covar_matrix[0, 0, 2].inverse().matmul(rhs_vector[0, 0, 2]).unsqueeze(0),
+                shifted_covar_matrix[0, 1, 0].inverse().matmul(rhs_vector[0, 1, 0]).unsqueeze(0),
+                shifted_covar_matrix[0, 1, 1].inverse().matmul(rhs_vector[0, 1, 1]).unsqueeze(0),
+                shifted_covar_matrix[0, 1, 2].inverse().matmul(rhs_vector[0, 1, 2]).unsqueeze(0),
+                shifted_covar_matrix[1, 0, 0].inverse().matmul(rhs_vector[1, 0, 0]).unsqueeze(0),
+                shifted_covar_matrix[1, 0, 1].inverse().matmul(rhs_vector[1, 0, 1]).unsqueeze(0),
+                shifted_covar_matrix[1, 0, 2].inverse().matmul(rhs_vector[1, 0, 2]).unsqueeze(0),
+                shifted_covar_matrix[1, 1, 0].inverse().matmul(rhs_vector[1, 1, 0]).unsqueeze(0),
+                shifted_covar_matrix[1, 1, 1].inverse().matmul(rhs_vector[1, 1, 1]).unsqueeze(0),
+                shifted_covar_matrix[1, 1, 2].inverse().matmul(rhs_vector[1, 1, 2]).unsqueeze(0),
+            ],
+            0,
+        ).view_as(rhs_vector)
+        scaled_inv_diag = (inv_scale / torch.ones(2, 3, 100)).unsqueeze(-1)
+        approx_solve = woodbury.woodbury_solve(
+            rhs_vector, piv_chol * scaled_inv_diag, woodbury_factor, scaled_inv_diag, inv_scale
+        )
+
+        self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
+
 
 if __name__ == "__main__":
     unittest.main()