Repository: cornellius-gp/gpytorch
Pull request #825 (2019-08-02 01:01:53)

Make tests less brittle w.r.t. upstream changes to warnings
Mitigates #824.

A proper solution should not count warnings, but check for presence/non-presence of the specific warning, see e.g. change in test_cholesky.py here.


Patch:
diff --git a/examples/README.md b/examples/README.md
index e3098e2d..8802c9f2 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -25,7 +25,7 @@ Some advanced techniques that you can apply to soup up these simple models:
 ## Specialty Models and Tasks
 
 - **Multitask GP Regression** - check out the examples in the [multitask GP folder](03_Multitask_GP_Regression)
-- **Bayesian Optimization** - check out the [BoTorch tutorials](https://botorch.org/tutorials/) or the [Ax tutorials](https://ax.dev/tutorials/), both of which use GPyTorch for BayesOpt.
+- **Bayesian Optimization** - example coming soon!
 
 
 ## Scalable GP Regression Models
diff --git a/gpytorch/lazy/__init__.py b/gpytorch/lazy/__init__.py
index ab94d04b..89e9f67f 100644
--- a/gpytorch/lazy/__init__.py
+++ b/gpytorch/lazy/__init__.py
@@ -5,7 +5,7 @@ from .added_diag_lazy_tensor import AddedDiagLazyTensor
 from .batch_repeat_lazy_tensor import BatchRepeatLazyTensor
 from .block_lazy_tensor import BlockLazyTensor
 from .block_diag_lazy_tensor import BlockDiagLazyTensor
-from .cached_cg_lazy_tensor import CachedCGLazyTensor
+from .cached_cg_lazy_tensor import CachedCGLazyTensor, ExtraComputationWarning
 from .cat_lazy_tensor import cat, CatLazyTensor
 from .chol_lazy_tensor import CholLazyTensor
 from .constant_mul_lazy_tensor import ConstantMulLazyTensor
@@ -39,6 +39,7 @@ __all__ = [
     "CholLazyTensor",
     "ConstantMulLazyTensor",
     "DiagLazyTensor",
+    "ExtraComputationWarning",
     "InterpolatedLazyTensor",
     "KroneckerProductLazyTensor",
     "MatmulLazyTensor",
diff --git a/gpytorch/lazy/cached_cg_lazy_tensor.py b/gpytorch/lazy/cached_cg_lazy_tensor.py
index 54951770..6295bdf5 100644
--- a/gpytorch/lazy/cached_cg_lazy_tensor.py
+++ b/gpytorch/lazy/cached_cg_lazy_tensor.py
@@ -7,6 +7,10 @@ from .chol_lazy_tensor import CholLazyTensor
 from .. import settings
 
 
+class ExtraComputationWarning(UserWarning):
+    pass
+
+
 class CachedCGLazyTensor(LazyTensor):
     """
     A LazyTensor wrapper that eagerly computes many CG calls in batch.
@@ -134,7 +138,8 @@ class CachedCGLazyTensor(LazyTensor):
         if settings.debug.on():
             warnings.warn(
                 "CachedCGLazyTensor had to run CG on a tensor of size {}. For best performance, this "
-                "LazyTensor should pre-register all vectors to run CG against.".format(rhs.shape)
+                "LazyTensor should pre-register all vectors to run CG against.".format(rhs.shape),
+                ExtraComputationWarning
             )
         return super(CachedCGLazyTensor, self)._cholesky_solve(rhs)
 
@@ -165,7 +170,8 @@ class CachedCGLazyTensor(LazyTensor):
             else:
                 if settings.debug.on():
                     warnings.warn(
-                        "CachedCGLazyTensor did not recognize the supplied probe vectors for tridiagonalization."
+                        "CachedCGLazyTensor did not recognize the supplied probe vectors for tridiagonalization.",
+                        ExtraComputationWarning
                     )
                 return super(CachedCGLazyTensor, self)._solve(rhs, preconditioner, num_tridiag=num_tridiag)
 
@@ -181,7 +187,8 @@ class CachedCGLazyTensor(LazyTensor):
         if settings.debug.on():
             warnings.warn(
                 "CachedCGLazyTensor had to run CG on a tensor of size {}. For best performance, this "
-                "LazyTensor should pre-register all vectors to run CG against.".format(rhs.shape)
+                "LazyTensor should pre-register all vectors to run CG against.".format(rhs.shape),
+                ExtraComputationWarning
             )
         return super(CachedCGLazyTensor, self)._solve(rhs, preconditioner, num_tridiag=num_tridiag)
 
@@ -245,3 +252,9 @@ class CachedCGLazyTensor(LazyTensor):
             logdet_term = self.base_lazy_tensor._chol_diag.pow(2).log().sum(-1)
 
         return inv_quad_term, logdet_term
+
+
+__all__ = [
+    "ExtraComputationWarning",
+    "CachedCGLazyTensor",
+]
diff --git a/gpytorch/variational/whitened_variational_strategy.py b/gpytorch/variational/whitened_variational_strategy.py
index 04defbd9..b77e4d32 100644
--- a/gpytorch/variational/whitened_variational_strategy.py
+++ b/gpytorch/variational/whitened_variational_strategy.py
@@ -189,8 +189,9 @@ class WhitenedVariationalStrategy(VariationalStrategy):
             if self.training:
                 data_covariance = DiagLazyTensor((data_data_covar.diag() - interp_data_data_var).clamp(0, math.inf))
             else:
-                neg_induc_data_data_covar = induc_induc_covar.inv_matmul(
-                    induc_data_covar, left_tensor=induc_data_covar.transpose(-1, -2).mul(-1)
+                neg_induc_data_data_covar = torch.matmul(
+                    induc_data_covar.transpose(-1, -2).mul(-1),
+                    induc_induc_covar.inv_matmul(induc_data_covar)
                 )
                 data_covariance = data_data_covar + neg_induc_data_data_covar
             predictive_covar = PsdSumLazyTensor(predictive_covar, data_covariance)


Test patch:
diff --git a/test/examples/test_svgp_gp_regression.py b/test/examples/test_svgp_gp_regression.py
index 17048977..9e8dec50 100644
--- a/test/examples/test_svgp_gp_regression.py
+++ b/test/examples/test_svgp_gp_regression.py
@@ -12,6 +12,7 @@ from gpytorch.models import AbstractVariationalGP
 from gpytorch.test.base_test_case import BaseTestCase
 from gpytorch.test.utils import least_used_cuda_device
 from gpytorch.variational import CholeskyVariationalDistribution, VariationalStrategy
+from gpytorch.lazy import ExtraComputationWarning
 from torch import optim
 
 
@@ -64,7 +65,7 @@ class TestSVGPRegression(BaseTestCase, unittest.TestCase):
         _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
         with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), gpytorch.settings.skip_logdet_forward(
             skip_logdet_forward
-        ), warnings.catch_warnings(record=True) as w, patch(
+        ), warnings.catch_warnings(record=True) as ws, patch(
             "gpytorch.utils.linear_cg", new=_wrapped_cg
         ) as linear_cg_mock:
             for _ in range(150):
@@ -74,26 +75,27 @@ class TestSVGPRegression(BaseTestCase, unittest.TestCase):
                 loss.backward()
                 optimizer.step()
 
+
+            for param in model.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+            for param in likelihood.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+
+            # Set back to eval mode
+            model.eval()
+            likelihood.eval()
+            test_preds = likelihood(model(train_x)).mean.squeeze()
+            mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
+            self.assertLess(mean_abs_error.item(), 1e-1)
+
             # Make sure CG was called (or not), and no warnings were thrown
-            self.assertEqual(len(w), 0)
             if cholesky:
                 self.assertFalse(linear_cg_mock.called)
             else:
                 self.assertTrue(linear_cg_mock.called)
-
-        for param in model.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-        for param in likelihood.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-
-        # Set back to eval mode
-        model.eval()
-        likelihood.eval()
-        test_preds = likelihood(model(train_x)).mean.squeeze()
-        mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
-        self.assertLess(mean_abs_error.item(), 1e-1)
+            self.assertFalse(any(issubclass(w.category, ExtraComputationWarning) for w in ws))
 
     def test_regression_error_skip_logdet_forward(self):
         return self.test_regression_error(skip_logdet_forward=True)
diff --git a/test/examples/test_whitened_svgp_gp_regression.py b/test/examples/test_whitened_svgp_gp_regression.py
index 8a8bf6b3..4fc6d58a 100644
--- a/test/examples/test_whitened_svgp_gp_regression.py
+++ b/test/examples/test_whitened_svgp_gp_regression.py
@@ -13,6 +13,7 @@ from gpytorch.models import AbstractVariationalGP
 from gpytorch.test.base_test_case import BaseTestCase
 from gpytorch.test.utils import least_used_cuda_device
 from gpytorch.variational import CholeskyVariationalDistribution, WhitenedVariationalStrategy
+from gpytorch.lazy import ExtraComputationWarning
 from torch import optim
 
 
@@ -65,7 +66,7 @@ class TestSVGPRegression(BaseTestCase, unittest.TestCase):
         _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
         with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), gpytorch.settings.skip_logdet_forward(
             skip_logdet_forward
-        ), warnings.catch_warnings(record=True) as w, patch(
+        ), warnings.catch_warnings(record=True) as ws, patch(
             "gpytorch.utils.linear_cg", new=_wrapped_cg
         ) as linear_cg_mock:
             for _ in range(200):
@@ -75,26 +76,27 @@ class TestSVGPRegression(BaseTestCase, unittest.TestCase):
                 loss.backward()
                 optimizer.step()
 
+
+            for param in model.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+            for param in likelihood.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+
+            # Set back to eval mode
+            model.eval()
+            likelihood.eval()
+            test_preds = likelihood(model(train_x)).mean.squeeze()
+            mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
+            self.assertLess(mean_abs_error.item(), 1e-1)
+
             # Make sure CG was called (or not), and no warnings were thrown
-            self.assertEqual(len(w), 0)
             if cholesky:
                 self.assertFalse(linear_cg_mock.called)
             else:
                 self.assertTrue(linear_cg_mock.called)
-
-        for param in model.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-        for param in likelihood.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-
-        # Set back to eval mode
-        model.eval()
-        likelihood.eval()
-        test_preds = likelihood(model(train_x)).mean.squeeze()
-        mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
-        self.assertLess(mean_abs_error.item(), 1e-1)
+            self.assertFalse(any(issubclass(w.category, ExtraComputationWarning) for w in ws))
 
     def test_regression_error_skip_logdet_forward(self):
         self.test_regression_error(skip_logdet_forward=True)
diff --git a/test/lazy/test_cached_cg_lazy_tensor.py b/test/lazy/test_cached_cg_lazy_tensor.py
index 849fa641..a84022b3 100644
--- a/test/lazy/test_cached_cg_lazy_tensor.py
+++ b/test/lazy/test_cached_cg_lazy_tensor.py
@@ -5,7 +5,7 @@ import torch
 import gpytorch
 import unittest
 import warnings
-from gpytorch.lazy import CachedCGLazyTensor, NonLazyTensor
+from gpytorch.lazy import CachedCGLazyTensor, NonLazyTensor, ExtraComputationWarning
 from gpytorch.utils.gradients import _ensure_symmetric_grad
 from gpytorch.test.lazy_tensor_test_case import LazyTensorTestCase
 from unittest.mock import MagicMock, patch
@@ -86,7 +86,7 @@ class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
         _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
         with patch("gpytorch.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
             with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), gpytorch.settings.cg_tolerance(1e-4):
-                with warnings.catch_warnings(record=True) as w:
+                with warnings.catch_warnings(record=True) as ws:
                     # Perform the inv_matmul
                     if lhs is not None:
                         res = lazy_tensor.inv_matmul(rhs, lhs)
@@ -95,9 +95,9 @@ class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
                         res = lazy_tensor.inv_matmul(rhs)
                         actual = evaluated.inverse().matmul(rhs_copy)
                     self.assertAllClose(res, actual, rtol=0.02, atol=1e-5)
-                    self.assertEqual(len(w), 0)
+                    self.assertFalse(any(issubclass(w.category, ExtraComputationWarning) for w in ws))
 
-                with warnings.catch_warnings(record=True) as w:
+                with warnings.catch_warnings(record=True) as ws:
                     # Perform backward pass
                     grad = torch.randn_like(res)
                     res.backward(gradient=grad)
@@ -112,11 +112,11 @@ class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
                     # Determine if we've called CG or not
                     # We shouldn't if we supplied a lhs
                     if lhs is None:
-                        self.assertEqual(len(w), 1)
+                        self.assertEqual(len([w for w in ws if issubclass(w.category, ExtraComputationWarning)]), 1)
                         if not cholesky and self.__class__.should_call_cg:
                             self.assertTrue(linear_cg_mock.called)
                     else:
-                        self.assertEqual(len(w), 0)
+                        self.assertFalse(any(issubclass(w.category, ExtraComputationWarning) for w in ws))
                         self.assertFalse(linear_cg_mock.called)
 
     def test_inv_matmul_vector(self):
@@ -156,9 +156,9 @@ class TestCachedCGLazyTensor(TestCachedCGLazyTensorNoLogdet):
         vecs = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
         vecs_copy = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
 
-        with gpytorch.settings.num_trace_samples(128), warnings.catch_warnings(record=True) as w:
+        with gpytorch.settings.num_trace_samples(128), warnings.catch_warnings(record=True) as ws:
             res_inv_quad, res_logdet = lazy_tensor.inv_quad_logdet(inv_quad_rhs=vecs, logdet=True)
-            self.assertEqual(len(w), 0)
+            self.assertFalse(any(issubclass(w.category, ExtraComputationWarning) for w in ws))
         res = res_inv_quad + res_logdet
 
         actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum(-2).sum(-1)
@@ -179,11 +179,11 @@ class TestCachedCGLazyTensor(TestCachedCGLazyTensorNoLogdet):
         vecs = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
         vecs_copy = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
 
-        with gpytorch.settings.num_trace_samples(128), warnings.catch_warnings(record=True) as w:
+        with gpytorch.settings.num_trace_samples(128), warnings.catch_warnings(record=True) as ws:
             res_inv_quad, res_logdet = lazy_tensor.inv_quad_logdet(
                 inv_quad_rhs=vecs, logdet=True, reduce_inv_quad=False
             )
-            self.assertEqual(len(w), 0)
+            self.assertFalse(any(issubclass(w.category, ExtraComputationWarning) for w in ws))
         res = res_inv_quad.sum(-1) + res_logdet
 
         actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum(-2).sum(-1)
diff --git a/test/utils/test_cholesky.py b/test/utils/test_cholesky.py
index d5839cc1..2056d031 100644
--- a/test/utils/test_cholesky.py
+++ b/test/utils/test_cholesky.py
@@ -66,9 +66,8 @@ class TestPSDSafeCholesky(unittest.TestCase):
                 L_exp = torch.cholesky(Aprime)
                 with warnings.catch_warnings(record=True) as w:
                     L_safe = psd_safe_cholesky(A)
-                    self.assertEqual(len(w), 1)
-                    self.assertTrue(issubclass(w[-1].category, RuntimeWarning))
-                    self.assertTrue("A not p.d., added jitter" in str(w[-1].message))
+                    self.assertTrue(any(issubclass(w_.category, RuntimeWarning) for w_ in w))
+                    self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
                 self.assertTrue(torch.allclose(L_exp, L_safe))
                 # user-defined value
                 Aprime = A.clone()
@@ -76,9 +75,8 @@ class TestPSDSafeCholesky(unittest.TestCase):
                 L_exp = torch.cholesky(Aprime)
                 with warnings.catch_warnings(record=True) as w:
                     L_safe = psd_safe_cholesky(A, jitter=1e-2)
-                    self.assertEqual(len(w), 1)
-                    self.assertTrue(issubclass(w[-1].category, RuntimeWarning))
-                    self.assertTrue("A not p.d., added jitter" in str(w[-1].message))
+                    self.assertTrue(any(issubclass(w_.category, RuntimeWarning) for w_ in w))
+                    self.assertTrue(any("A not p.d., added jitter" in str(w_.message) for w_ in w))
                 self.assertTrue(torch.allclose(L_exp, L_safe))
 
     def test_psd_safe_cholesky_psd_cuda(self, cuda=False):