Repository: cornellius-gp/gpytorch
Pull request #500 (2019-02-05 22:53:30)

Wrapper around cholesky to automatically add jitter if matrix is not p.d.
`psd_safe_cholesky` is a simple wrapper around `torch.cholesky` that adds jitter if the matrix is not p.d.. If computing `torch.cholesky` on that still fails a `RuntimeError` is raised. This also hackfixes the  newly discovered pytorch bug: https://github.com/pytorch/pytorch/issues/16780

Balandat: Some context: This is nice to have for optimizing q-batch acquisition functions in BayesOpt, where as the solution improves we end up with points close to each other, which can lead to degenerate posteriors. This currently either errors out or silently produces NaNs on the GPU (see bug).
Balandat: Not sure about the test failures here things ran fine locally - I probably messed something up staging this.
Balandat: hmm, why doesn't this run Travis on the recent commit?

Patch:
diff --git a/gpytorch/lazy/block_diag_lazy_tensor.py b/gpytorch/lazy/block_diag_lazy_tensor.py
index e128e642..75d92e3a 100644
--- a/gpytorch/lazy/block_diag_lazy_tensor.py
+++ b/gpytorch/lazy/block_diag_lazy_tensor.py
@@ -3,6 +3,7 @@
 import torch
 
 from .. import settings
+from ..utils.cholesky import psd_safe_cholesky
 from ..utils.memoize import cached
 from .block_lazy_tensor import BlockLazyTensor
 from .non_lazy_tensor import NonLazyTensor
@@ -127,8 +128,8 @@ class BlockDiagLazyTensor(BlockLazyTensor):
         if settings.fast_computations.covar_root_decomposition.on():
             res = self.__class__(self.base_lazy_tensor.root_decomposition().root, num_blocks=self.num_blocks)
         else:
-            chol = torch.cholesky(self.base_lazy_tensor.evaluate())
-            res = self.__class__(NonLazyTensor(chol), num_blocks=self.num_blocks)
+            chol = psd_safe_cholesky(self.base_lazy_tensor.evaluate())
+            res = self.__class__(NonLazyTensor(chol), block_dim=self.block_dim)
         return RootLazyTensor(res)
 
     def zero_mean_mvn_samples(self, num_samples):
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 181e4d64..cf1f447d 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -12,6 +12,7 @@ from ..functions._matmul import Matmul
 from ..functions._root_decomposition import RootDecomposition
 from ..utils import linear_cg
 from ..utils.broadcasting import _matmul_broadcast_shape
+from ..utils.cholesky import psd_safe_cholesky
 from ..utils.deprecation import _deprecate_renamed_methods
 from ..utils.memoize import cached
 from ..utils.qr import batch_qr
@@ -1021,22 +1022,16 @@ class LazyTensor(object):
 
         if (self.matrix_shape.numel() <= settings.max_cholesky_numel.value()
                 or settings.fast_computations.covar_root_decomposition.off()):
-            try:
-                res = torch.cholesky(self.evaluate())
-                return RootLazyTensor(res)
-            except RuntimeError as e:
-                warnings.warn(
-                    "Runtime Error when computing Cholesky decomposition: {}. Using RootDecomposition.".format(e)
-                )
-
-        res, _ = RootDecomposition(
-            self.representation_tree(),
-            max_iter=self.root_decomposition_size(),
-            dtype=self.dtype,
-            device=self.device,
-            batch_shape=self.batch_shape,
-            matrix_shape=self.matrix_shape,
-        )(*self.representation())
+            res = psd_safe_cholesky(self.evaluate())
+        else:
+            res, _ = RootDecomposition(
+                self.representation_tree(),
+                max_iter=self.root_decomposition_size(),
+                dtype=self.dtype,
+                device=self.device,
+                batch_shape=self.batch_shape,
+                matrix_shape=self.matrix_shape,
+            )(*self.representation())
 
         return RootLazyTensor(res)
 
diff --git a/gpytorch/utils/cholesky.py b/gpytorch/utils/cholesky.py
index f53ec1bf..d598ab1a 100644
--- a/gpytorch/utils/cholesky.py
+++ b/gpytorch/utils/cholesky.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import warnings
 
 import torch
 
@@ -82,3 +83,43 @@ def tridiag_batch_potrs(tensor, chol_trid, upper=True):
         solution[:, i, :].copy_(inner_part / chol_main_diag[:, i].unsqueeze(-1))
 
     return solution
+
+
+def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
+    """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
+
+    Args:
+        :attr:`A` (Tensor):
+            The tensor to compute the Cholesky decomposition of
+        :attr:`upper` (bool, optional):
+            See torch.cholesky
+        :attr:`out` (Tensor, optional):
+            See torch.cholesky
+        :attr:`jitter` (float, optional):
+            The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen as 1e-5 (float) or 1e-7 (double)
+    """
+    try:
+        L = torch.cholesky(A, upper=upper, out=out)
+        # TODO: Remove once fixed in pytorch (#16780)
+        if A.dim() > 2 and A.is_cuda:
+            if torch.isnan(L if out is None else out).any():
+                raise RuntimeError
+    except RuntimeError:
+        if jitter is None:
+            jitter = 1e-5 if A.dtype == torch.float32 else 1e-7
+        idx = torch.arange(A.shape[-1], device=A.device)
+        Aprime = A.clone()
+        Aprime[..., idx, idx] += jitter
+        try:
+            L = torch.cholesky(Aprime, upper=upper, out=out)
+            # TODO: Remove once fixed in pytorch (#16780)
+            if A.dim() > 2 and A.is_cuda:
+                if torch.isnan(L if out is None else out).any():
+                    raise RuntimeError("singular")
+        except RuntimeError as e:
+            if "singular" in e.args[0]:
+                raise RuntimeError("Adding jitter of {} to the diagonal did not make A p.d.".format(jitter))
+        warnings.warn("A not p.d., added jitter of {} to the diagonal".format(jitter), RuntimeWarning)
+
+    if out is None:
+        return L


Test patch:
diff --git a/test/utils/test_cholesky.py b/test/utils/test_cholesky.py
index 3aa466ea..3631e2c0 100644
--- a/test/utils/test_cholesky.py
+++ b/test/utils/test_cholesky.py
@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
 
-import torch
 import unittest
+import warnings
 from test._utils import approx_equal
-from gpytorch.utils.cholesky import tridiag_batch_potrf, tridiag_batch_potrs
+
+import torch
+from gpytorch.utils.cholesky import psd_safe_cholesky, tridiag_batch_potrf, tridiag_batch_potrs
 
 
 class TestTriDiag(unittest.TestCase):
@@ -22,5 +24,80 @@ class TestTriDiag(unittest.TestCase):
         )
 
 
+class TestPSDSafeCholesky(unittest.TestCase):
+    def _gen_test_psd(self):
+        return torch.tensor([[[0.25, -0.75], [-0.75, 2.25]], [[1.0, 1.0], [1.0, 1.0]]])
+
+    def test_psd_safe_cholesky_pd(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            for batch_mode in (False, True):
+                if batch_mode:
+                    A = self._gen_test_psd().to(device=device, dtype=dtype)
+                    D = torch.eye(2).type_as(A).unsqueeze(0).repeat(2, 1, 1)
+                else:
+                    A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
+                    D = torch.eye(2).type_as(A)
+                A += D
+                # basic
+                L = torch.cholesky(A)
+                L_safe = psd_safe_cholesky(A)
+                self.assertTrue(torch.allclose(L, L_safe))
+                # upper
+                L = torch.cholesky(A, upper=True)
+                L_safe = psd_safe_cholesky(A, upper=True)
+                self.assertTrue(torch.allclose(L, L_safe))
+                # output tensors
+                L = torch.empty_like(A)
+                L_safe = torch.empty_like(A)
+                torch.cholesky(A, out=L)
+                psd_safe_cholesky(A, out=L_safe)
+                self.assertTrue(torch.allclose(L, L_safe))
+                # output tensors, upper
+                torch.cholesky(A, upper=True, out=L)
+                psd_safe_cholesky(A, upper=True, out=L_safe)
+                self.assertTrue(torch.allclose(L, L_safe))
+                # make sure jitter doesn't do anything if p.d.
+                L = torch.cholesky(A)
+                L_safe = psd_safe_cholesky(A, jitter=1e-2)
+                self.assertTrue(torch.allclose(L, L_safe))
+
+    def test_psd_safe_cholesky_pd_cuda(self, cuda=False):
+        if torch.cuda.is_available():
+            self.test_psd_safe_cholesky_pd(cuda=True)
+
+    def test_psd_safe_cholesky_psd(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        for dtype in (torch.float, torch.double):
+            for batch_mode in (False, True):
+                if batch_mode:
+                    A = self._gen_test_psd().to(device=device, dtype=dtype)
+                else:
+                    A = self._gen_test_psd()[0].to(device=device, dtype=dtype)
+                idx = torch.arange(A.shape[-1], device=A.device)
+                # default values
+                Aprime = A.clone()
+                Aprime[..., idx, idx] += 1e-5 if A.dtype == torch.float32 else 1e-7
+                L_exp = torch.cholesky(Aprime)
+                with warnings.catch_warnings(record=True) as ws:
+                    L_safe = psd_safe_cholesky(A)
+                    self.assertEqual(len(ws), 1)
+                    self.assertEqual(ws[-1].category, RuntimeWarning)
+                self.assertTrue(torch.allclose(L_exp, L_safe))
+                # user-defined value
+                Aprime = A.clone()
+                Aprime[..., idx, idx] += 1e-2
+                L_exp = torch.cholesky(Aprime)
+                with warnings.catch_warnings(record=True) as ws:
+                    L_safe = psd_safe_cholesky(A, jitter=1e-2)
+                    self.assertEqual(len(ws), 1)
+                    self.assertEqual(ws[-1].category, RuntimeWarning)
+                self.assertTrue(torch.allclose(L_exp, L_safe))
+
+    def test_psd_safe_cholesky_psd_cuda(self, cuda=False):
+        if torch.cuda.is_available():
+            self.test_psd_safe_cholesky_psd(cuda=True)
+
+
 if __name__ == "__main__":
     unittest.main()