catch runtime errors with ill-conditioned covar (#1095)

Summary:
Pull Request resolved: #1095

see title.

Reviewed By: dme65, Balandat

Differential Revision: D34424978

fbshipit-source-id: 59f90a26b83dd9071b4515f2dc97ede8b2568dab

Author: sdaulton

botorch/utils/low_rank.py

@@ -16,7 +16,7 @@
 from gpytorch.lazy import BlockDiagLazyTensor
 from gpytorch.lazy.lazy_tensor import LazyTensor
 from gpytorch.utils.cholesky import psd_safe_cholesky
-from gpytorch.utils.errors import NanError
+from gpytorch.utils.errors import NotPSDError, NanError
 from torch import Tensor
 
 
@@ -121,9 +121,16 @@ def sample_cached_cholesky(
     # and bl_chol := x^T
     # bl_chol is the new `(batch_shape) x q x n`-dim bottom left block
     # of the cholesky decomposition
-    bl_chol = torch.triangular_solve(
-        bl.transpose(-2, -1), baseline_L, upper=False
-    ).solution.transpose(-2, -1)
+    # TODO: remove the exception handling, when the pytorch
+    # version requirement is bumped to >= 1.10
+    try:
+        bl_chol = torch.triangular_solve(
+            bl.transpose(-2, -1), baseline_L, upper=False
+        ).solution.transpose(-2, -1)
+    except RuntimeError as e:
+        if "singular" in str(e):
+            raise NotPSDError(f"triangular_solve failed with RuntimeError: {e}")
+        raise e
     # Compute the new bottom right block of the Cholesky
     # decomposition via:
     # Cholesky(K(X, X) - bl_chol @ bl_chol^T)

test/utils/test_low_rank.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+from unittest import mock
+
 import torch
 from botorch.exceptions.errors import BotorchError
 from botorch.models.gp_regression import SingleTaskGP
@@ -16,7 +18,7 @@
 )
 from gpytorch.lazy import lazify
 from gpytorch.lazy.block_diag_lazy_tensor import BlockDiagLazyTensor
-from gpytorch.utils.errors import NanError
+from gpytorch.utils.errors import NotPSDError, NanError
 
 
 class TestExtractBatchCovar(BotorchTestCase):
@@ -197,3 +199,32 @@ def test_sample_cached_cholesky(self):
                                     base_samples=sampler.base_samples.detach().clone(),
                                     sample_shape=sampler.sample_shape,
                                 )
+                            # test triangular solve raising RuntimeError
+                            test_posterior.mvn.loc = torch.full_like(
+                                test_posterior.mvn.loc, 0.0
+                            )
+                            base_samples = sampler.base_samples.detach().clone()
+                            with mock.patch(
+                                "botorch.utils.low_rank.torch.triangular_solve",
+                                side_effect=RuntimeError("singular"),
+                            ):
+                                with self.assertRaises(NotPSDError):
+                                    sample_cached_cholesky(
+                                        posterior=test_posterior,
+                                        baseline_L=baseline_L,
+                                        q=q,
+                                        base_samples=base_samples,
+                                        sample_shape=sampler.sample_shape,
+                                    )
+                            with mock.patch(
+                                "botorch.utils.low_rank.torch.triangular_solve",
+                                side_effect=RuntimeError(""),
+                            ):
+                                with self.assertRaises(RuntimeError):
+                                    sample_cached_cholesky(
+                                        posterior=test_posterior,
+                                        baseline_L=baseline_L,
+                                        q=q,
+                                        base_samples=base_samples,
+                                        sample_shape=sampler.sample_shape,
+                                    )