Repository: cornellius-gp/gpytorch
Pull request #506 (2019-02-07 16:56:38)

Add psd_safe_cholesky back in
Fx rebasing issue. Never perform Choleksy for MulLazyTensor.


Patch:
diff --git a/gpytorch/lazy/block_diag_lazy_tensor.py b/gpytorch/lazy/block_diag_lazy_tensor.py
index e128e642..b6a90001 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,7 +128,7 @@ 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())
+            chol = psd_safe_cholesky(self.base_lazy_tensor.evaluate())
             res = self.__class__(NonLazyTensor(chol), num_blocks=self.num_blocks)
         return RootLazyTensor(res)
 
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 181e4d64..3d85bd13 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -2,6 +2,7 @@
 
 import math
 import warnings
+
 import gpytorch
 import torch
 
@@ -12,6 +13,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
@@ -103,7 +105,7 @@ class LazyTensor(object):
             n_tridiag=num_tridiag,
             max_iter=settings.max_cg_iterations.value(),
             max_tridiag_iter=settings.max_lanczos_quadrature_iterations.value(),
-            preconditioner=preconditioner
+            preconditioner=preconditioner,
         )
 
     def _size(self):
@@ -678,8 +680,9 @@ class LazyTensor(object):
                 )
 
         func = InvMatmul(
-            self.representation_tree(), preconditioner=self._inv_matmul_preconditioner(),
-            has_left=(left_tensor is not None)
+            self.representation_tree(),
+            preconditioner=self._inv_matmul_preconditioner(),
+            has_left=(left_tensor is not None),
         )
         if left_tensor is None:
             return func(right_tensor, *self.representation())
@@ -1013,17 +1016,20 @@ class LazyTensor(object):
         low-rank version of a matrix
         """
         from .root_lazy_tensor import RootLazyTensor
+        from .mul_lazy_tensor import MulLazyTensor
+
         if not self.is_square:
             raise RuntimeError(
                 "root_decomposition only operates on (batches of) square (symmetric) LazyTensors. "
                 "Got a {} of size {}.".format(self.__class__.__name__, self.size())
             )
 
-        if (self.matrix_shape.numel() <= settings.max_cholesky_numel.value()
-                or settings.fast_computations.covar_root_decomposition.off()):
+        if not isinstance(self, MulLazyTensor) and (
+            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)
+                return RootLazyTensor(psd_safe_cholesky(self.evaluate()))
             except RuntimeError as e:
                 warnings.warn(
                     "Runtime Error when computing Cholesky decomposition: {}. Using RootDecomposition.".format(e)
@@ -1037,7 +1043,6 @@ class LazyTensor(object):
             batch_shape=self.batch_shape,
             matrix_shape=self.matrix_shape,
         )(*self.representation())
-
         return RootLazyTensor(res)
 
     @cached
@@ -1404,13 +1409,6 @@ def delazify(obj):
         raise TypeError("object of class {} cannot be made into a Tensor".format(obj.__class__.__name__))
 
 
-_deprecate_renamed_methods(
-    LazyTensor,
-    inv_quad_log_det="inv_quad_logdet",
-    log_det="logdet",
-)
+_deprecate_renamed_methods(LazyTensor, inv_quad_log_det="inv_quad_logdet", log_det="logdet")
 
-__all__ = [
-    "LazyTensor",
-    "delazify",
-]
+__all__ = ["LazyTensor", "delazify"]
diff --git a/gpytorch/utils/cholesky.py b/gpytorch/utils/cholesky.py
index f53ec1bf..66d84e44 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-6 (float) or 1e-8 (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-6 if A.dtype == torch.float32 else 1e-8
+        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..f349cd88 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-6 if A.dtype == torch.float32 else 1e-8
+                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()