Repository: cornellius-gp/gpytorch
Pull request #930 (2019-11-08 20:41:18)

Add preconditioner override option for AddedDiagLazyTensor
Adding a `preconditioner_override` option that allows specification of arbitrary preconditioners for `AddedDiagLazyTensor`. Example usage is given in the associated unit test but broadly looks like:

```
 overrode_lt = AddedDiagLazyTensor(
       RootLazyTensor(tensor),
       DiagLazyTensor(diag),
       preconditioner_override=nonstandard_preconditioner,
 )
```

The original preconditioner is still left in the class, but renamed as `_default_preconditioner`

Potential use cases include easier testing of new preconditioners as well as data driven preconditioners such as recycled subspaces from old solves.


Patch:
diff --git a/gpytorch/lazy/added_diag_lazy_tensor.py b/gpytorch/lazy/added_diag_lazy_tensor.py
index 49903d7e..4ae6da04 100644
--- a/gpytorch/lazy/added_diag_lazy_tensor.py
+++ b/gpytorch/lazy/added_diag_lazy_tensor.py
@@ -17,16 +17,23 @@ class AddedDiagLazyTensor(SumLazyTensor):
     a DiagLazyTensor.
     """
 
-    def __init__(self, *lazy_tensors):
+    def __init__(self, *lazy_tensors, preconditioner_override=None):
         lazy_tensors = list(lazy_tensors)
-        super(AddedDiagLazyTensor, self).__init__(*lazy_tensors)
+        super(AddedDiagLazyTensor, self).__init__(
+            *lazy_tensors, preconditioner_override=preconditioner_override
+        )
         if len(lazy_tensors) > 2:
             raise RuntimeError("An AddedDiagLazyTensor can only have two components")
 
         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.")
+        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."
+            )
         elif isinstance(lazy_tensors[0], DiagLazyTensor):
             self._diag_tensor = lazy_tensors[0]
             self._lazy_tensor = lazy_tensors[1]
@@ -34,13 +41,21 @@ class AddedDiagLazyTensor(SumLazyTensor):
             self._diag_tensor = lazy_tensors[1]
             self._lazy_tensor = lazy_tensors[0]
         else:
-            raise RuntimeError("One of the LazyTensors input to AddedDiagLazyTensor must be a DiagLazyTensor!")
+            raise RuntimeError(
+                "One of the LazyTensors input to AddedDiagLazyTensor must be a DiagLazyTensor!"
+            )
+
+        self.preconditioner_override = preconditioner_override
 
     def _matmul(self, rhs):
-        return torch.addcmul(self._lazy_tensor._matmul(rhs), self._diag_tensor._diag.unsqueeze(-1), rhs)
+        return torch.addcmul(
+            self._lazy_tensor._matmul(rhs), self._diag_tensor._diag.unsqueeze(-1), rhs
+        )
 
     def add_diag(self, added_diag):
-        return AddedDiagLazyTensor(self._lazy_tensor, self._diag_tensor.add_diag(added_diag))
+        return AddedDiagLazyTensor(
+            self._lazy_tensor, self._diag_tensor.add_diag(added_diag)
+        )
 
     def __add__(self, other):
         from .diag_lazy_tensor import DiagLazyTensor
@@ -51,25 +66,39 @@ class AddedDiagLazyTensor(SumLazyTensor):
             return AddedDiagLazyTensor(self._lazy_tensor + other, self._diag_tensor)
 
     def _preconditioner(self):
-        if settings.max_preconditioner_size.value() == 0 or self.size(-1) < settings.min_preconditioning_size.value():
+        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 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)
+            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."
                 )
                 return None, None, None
 
-            if self._piv_chol_self.dim() == 2:  # TODO: Whenever PyTorch supports batch mode
+            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)
 
-                self.constant_diag = torch.equal(self._noise, self._noise[0] * torch.ones_like(self._noise))
+                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:
                     # We can factor out the noise for for both QR and solves.
@@ -80,28 +109,53 @@ class AddedDiagLazyTensor(SumLazyTensor):
                     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 = (
+                        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_logdet_cache = (
+                        logdet.view(*batch_shape)
+                        if len(batch_shape)
+                        else logdet.squeeze()
+                    )
 
                 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._r_cache = torch.qr(
+                        torch.cat((self._piv_chol_self / self._noise.sqrt(), eye))
+                    )
                     self._q_cache = self._q_cache[:n, :] / self._noise.sqrt()
 
-                    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()
+                    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()
+                    )
 
             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._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._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)
+            self.preconditioner_lt = PsdSumLazyTensor(
+                RootLazyTensor(self._piv_chol_self), self._diag_tensor
+            )
 
         # NOTE to future self:
         # We cannot memoize this precondition closure
@@ -109,10 +163,14 @@ class AddedDiagLazyTensor(SumLazyTensor):
         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)))
+                    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))
+                    return (tensor / self._noise) - self._q_cache.matmul(
+                        self._q_cache.t().matmul(tensor)
+                    )
 
             else:
                 res = woodbury.woodbury_solve(
@@ -124,4 +182,4 @@ class AddedDiagLazyTensor(SumLazyTensor):
                 )
                 return res
 
-        return precondition_closure, self.preconditioner_lt, self._precond_logdet_cache
+        return (precondition_closure, self.preconditioner_lt, self._precond_logdet_cache)
diff --git a/gpytorch/lazy/sum_lazy_tensor.py b/gpytorch/lazy/sum_lazy_tensor.py
index ee49a653..c97e746d 100644
--- a/gpytorch/lazy/sum_lazy_tensor.py
+++ b/gpytorch/lazy/sum_lazy_tensor.py
@@ -7,14 +7,14 @@ from .zero_lazy_tensor import ZeroLazyTensor
 
 
 class SumLazyTensor(LazyTensor):
-    def __init__(self, *lazy_tensors):
+    def __init__(self, *lazy_tensors, **kwargs):
         lazy_tensors = list(lazy_tensors)
         for i, lazy_tensor in enumerate(lazy_tensors):
             try:
                 lazy_tensors[i] = lazify(lazy_tensor)
             except TypeError:
                 raise TypeError("All arguments of a SumLazyTensor should be LazyTensors or Tensors")
-        super(SumLazyTensor, self).__init__(*lazy_tensors)
+        super(SumLazyTensor, self).__init__(*lazy_tensors, **kwargs)
 
         self.lazy_tensors = lazy_tensors
 


Test patch:
diff --git a/test/lazy/test_added_diag_lazy_tensor.py b/test/lazy/test_added_diag_lazy_tensor.py
index a4182fb4..45f73f60 100644
--- a/test/lazy/test_added_diag_lazy_tensor.py
+++ b/test/lazy/test_added_diag_lazy_tensor.py
@@ -2,7 +2,12 @@
 
 import torch
 import unittest
-from gpytorch.lazy import NonLazyTensor, DiagLazyTensor, AddedDiagLazyTensor
+from gpytorch.lazy import (
+    NonLazyTensor,
+    DiagLazyTensor,
+    AddedDiagLazyTensor,
+    RootLazyTensor,
+)
 from gpytorch.test.lazy_tensor_test_case import LazyTensorTestCase
 
 
@@ -30,7 +35,12 @@ class TestAddedDiagLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
         tensor = torch.randn(3, 5, 5)
         tensor = tensor.transpose(-1, -2).matmul(tensor).detach()
         diag = torch.tensor(
-            [[1.0, 2.0, 4.0, 2.0, 3.0], [2.0, 1.0, 2.0, 1.0, 4.0], [1.0, 2.0, 2.0, 3.0, 4.0]], requires_grad=True
+            [
+                [1.0, 2.0, 4.0, 2.0, 3.0],
+                [2.0, 1.0, 2.0, 1.0, 4.0],
+                [1.0, 2.0, 2.0, 3.0, 4.0],
+            ],
+            requires_grad=True,
         )
         return AddedDiagLazyTensor(NonLazyTensor(tensor), DiagLazyTensor(diag))
 
@@ -51,7 +61,12 @@ class TestAddedDiagLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
         tensor = tensor.transpose(-1, -2).matmul(tensor).detach()
         diag = (
             torch.tensor(
-                [[1.0, 2.0, 4.0, 2.0, 3.0], [2.0, 1.0, 2.0, 1.0, 4.0], [1.0, 2.0, 2.0, 3.0, 4.0]], requires_grad=True
+                [
+                    [1.0, 2.0, 4.0, 2.0, 3.0],
+                    [2.0, 1.0, 2.0, 1.0, 4.0],
+                    [1.0, 2.0, 2.0, 3.0, 4.0],
+                ],
+                requires_grad=True,
             )
             .repeat(4, 1, 1)
             .detach()
@@ -64,5 +79,48 @@ class TestAddedDiagLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
         return tensor + torch.diag_embed(diag, dim1=-2, dim2=-1)
 
 
+class TestAddedDiagLazyTensorPrecondOverride(unittest.TestCase):
+    def test_precond_solve(self):
+        seed = 4
+        torch.random.manual_seed(seed)
+
+        tensor = torch.randn(1000, 800)
+        diag = torch.abs(torch.randn(1000))
+
+        standard_lt = AddedDiagLazyTensor(RootLazyTensor(tensor), DiagLazyTensor(diag))
+        evals, evecs = torch.symeig(standard_lt.evaluate(), eigenvectors=True)
+
+        # this preconditioner is a simple example of near deflation
+        def nonstandard_preconditioner(self):
+            top_100_evecs = evecs[:, :100]
+            top_100_evals = evals[:100] + 0.2 * torch.randn(100)
+
+            precond_lt = RootLazyTensor(top_100_evecs @ torch.diag(top_100_evals ** 0.5))
+            logdet = top_100_evals.log().sum()
+
+            def precond_closure(rhs):
+                rhs2 = top_100_evecs.t() @ rhs
+                return top_100_evecs @ torch.diag(1.0 / top_100_evals) @ rhs2
+
+            return precond_closure, precond_lt, logdet
+
+        overrode_lt = AddedDiagLazyTensor(
+            RootLazyTensor(tensor),
+            DiagLazyTensor(diag),
+            preconditioner_override=nonstandard_preconditioner,
+        )
+
+        # compute a solve - mostly to make sure that we can actually perform the solve
+        rhs = torch.randn(1000, 1)
+        standard_solve = standard_lt.inv_matmul(rhs)
+        overrode_solve = overrode_lt.inv_matmul(rhs)
+
+        # gut checking that our preconditioner is not breaking anything
+        self.assertEqual(standard_solve.shape, overrode_solve.shape)
+        self.assertLess(
+            torch.norm(standard_solve - overrode_solve) / standard_solve.norm(), 1.0
+        )
+
+
 if __name__ == "__main__":
     unittest.main()