Merge branch 'stable_qr_fix' of https://github.com/wjmaddox/gpytorch into stable_qr_fix

Author: wjmaddox

gpytorch/kernels/kernel.py

@@ -331,8 +331,8 @@ def covar_dist(
         return res
 
     def named_sub_kernels(self):
-        for name, module in self._modules.items():
-            if isinstance(module, Kernel):
+        for name, module in self.named_modules():
+            if module is not self and isinstance(module, Kernel):
                 yield name, module
 
     def num_outputs_per_input(self, x1, x2):

gpytorch/lazy/lazy_tensor.py

@@ -780,9 +780,9 @@ def cat_rows(self, cross_mat, new_mat, generate_roots=True, **root_decomp_kwargs
             A = self
 
         # form matrix C = [A B; B^T D], where A = self, B = cross_mat, D = new_mat
-        upper_row = CatLazyTensor(A, B, dim=-2)
-        lower_row = CatLazyTensor(B.transpose(-1, -2), D, dim=-2)
-        new_lazy_tensor = CatLazyTensor(upper_row, lower_row, dim=-1)
+        upper_row = CatLazyTensor(A, B, dim=-2, output_device=A.device)
+        lower_row = CatLazyTensor(B.transpose(-1, -2), D, dim=-2, output_device=A.device)
+        new_lazy_tensor = CatLazyTensor(upper_row, lower_row, dim=-1, output_device=A.device)
 
         # if the old lazy tensor does not have either a root decomposition or a root inverse decomposition
         # don't create one

gpytorch/lazy/low_rank_root_added_diag_lazy_tensor.py

@@ -124,3 +124,32 @@ def inv_quad_logdet(self, inv_quad_rhs=None, logdet=False, reduce_inv_quad=True)
             logdet_term = self._logdet()
 
         return inv_quad_term, logdet_term
+
+    def inv_matmul(self, right_tensor, left_tensor=None):
+        if not self.is_square:
+            raise RuntimeError(
+                "inv_matmul only operates on (batches of) square (positive semi-definite) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
+        if self.dim() == 2 and right_tensor.dim() == 1:
+            if self.shape[-1] != right_tensor.numel():
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                        self.shape, right_tensor.shape
+                    )
+                )
+
+        squeeze_solve = False
+        if right_tensor.ndimension() == 1:
+            right_tensor = right_tensor.unsqueeze(-1)
+            squeeze_solve = True
+
+        solve = self._solve(right_tensor)
+        if squeeze_solve:
+            solve = solve.squeeze(-1)
+
+        if left_tensor is not None:
+            return left_tensor @ solve
+        else:
+            return solve

test/examples/test_sgpr_regression.py

@@ -3,6 +3,7 @@
 import os
 import random
 import unittest
+from unittest.mock import MagicMock, patch
 import warnings
 from math import exp, pi
 
@@ -14,6 +15,7 @@
 from gpytorch.means import ConstantMean
 from gpytorch.priors import SmoothedBoxPrior
 from gpytorch.test.utils import least_used_cuda_device
+from gpytorch.utils.cholesky import CHOLESKY_METHOD
 from gpytorch.utils.warnings import NumericalWarning
 from torch import optim
 
@@ -65,41 +67,56 @@ def tearDown(self):
         if hasattr(self, "rng_state"):
             torch.set_rng_state(self.rng_state)
 
-    def test_sgpr_mean_abs_error(self):
+    def test_sgpr_mean_abs_error(self, cuda=False):
         # Suppress numerical warnings
         warnings.simplefilter("ignore", NumericalWarning)
 
-        train_x, train_y, test_x, test_y = make_data()
+        train_x, train_y, test_x, test_y = make_data(cuda=cuda)
         likelihood = GaussianLikelihood()
         gp_model = GPRegressionModel(train_x, train_y, likelihood)
         mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)
 
-        # Optimize the model
-        gp_model.train()
-        likelihood.train()
+        if cuda:
+            gp_model = gp_model.cuda()
+            likelihood = likelihood.cuda()
 
-        optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
-        for _ in range(30):
-            optimizer.zero_grad()
-            output = gp_model(train_x)
-            loss = -mll(output, train_y)
-            loss.backward()
-            optimizer.step()
+        # Mock cholesky
+        _wrapped_cholesky = MagicMock(
+            wraps=torch.linalg.cholesky if CHOLESKY_METHOD == "torch.linalg.cholesky" else torch.linalg.cholesky_ex
+        )
+        with patch(CHOLESKY_METHOD, new=_wrapped_cholesky) as cholesky_mock:
 
-            # Check that we have the right LazyTensor type
-            kernel = likelihood(gp_model(train_x)).lazy_covariance_matrix.evaluate_kernel()
-            self.assertIsInstance(kernel, gpytorch.lazy.LowRankRootAddedDiagLazyTensor)
+            # Optimize the model
+            gp_model.train()
+            likelihood.train()
 
-        for param in gp_model.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
+            optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
+            for _ in range(30):
+                optimizer.zero_grad()
+                output = gp_model(train_x)
+                loss = -mll(output, train_y)
+                loss.backward()
+                optimizer.step()
 
-        # Test the model
-        gp_model.eval()
-        likelihood.eval()
+                # Check that we have the right LazyTensor type
+                kernel = likelihood(gp_model(train_x)).lazy_covariance_matrix.evaluate_kernel()
+                self.assertIsInstance(kernel, gpytorch.lazy.LowRankRootAddedDiagLazyTensor)
 
-        test_preds = likelihood(gp_model(test_x)).mean
-        mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
+            for param in gp_model.parameters():
+                self.assertTrue(param.grad is not None)
+                self.assertGreater(param.grad.norm().item(), 0)
+
+            # Test the model
+            gp_model.eval()
+            likelihood.eval()
+
+            test_preds = likelihood(gp_model(test_x)).mean
+            mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
+            cholesky_mock.assert_called()  # We SHOULD call Cholesky...
+            for chol_arg in cholesky_mock.call_args_list:
+                first_arg = chol_arg[0][0]
+                self.assertTrue(torch.is_tensor(first_arg))
+                self.assertTrue(first_arg.size(-1) == gp_model.covar_module.inducing_points.size(-2))
 
         self.assertLess(mean_abs_error.squeeze().item(), 0.1)
 
@@ -123,62 +140,9 @@ def test_sgpr_mean_abs_error_cuda(self):
 
         if not torch.cuda.is_available():
             return
-        with least_used_cuda_device():
-            train_x, train_y, test_x, test_y = make_data(cuda=True)
-            likelihood = GaussianLikelihood().cuda()
-            gp_model = GPRegressionModel(train_x, train_y, likelihood).cuda()
-            mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)
-
-            # Test the model before optimization
-            gp_model.eval()
-            likelihood.eval()
-            test_preds = likelihood(gp_model(test_x)).mean
-            mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
-            self.assertLess(mean_abs_error.squeeze().item(), 0.02)
-
-            # Test variances before optimization
-            test_vars = likelihood(gp_model(test_x)).variance
-            self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
-            self.assertGreater(test_vars.min().item() + 0.1, likelihood.noise.item())
-            self.assertLess(
-                test_vars.max().item() - 0.05,
-                likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
-            )
-
-            # Optimize the model
-            gp_model.train()
-            likelihood.train()
-
-            optimizer = optim.Adam(gp_model.parameters(), lr=0.1)
-            optimizer.n_iter = 0
-            for _ in range(25):
-                optimizer.zero_grad()
-                output = gp_model(train_x)
-                loss = -mll(output, train_y)
-                loss.backward()
-                optimizer.n_iter += 1
-                optimizer.step()
 
-            for param in gp_model.parameters():
-                self.assertTrue(param.grad is not None)
-                self.assertGreater(param.grad.norm().item(), 0)
-
-            # Test the model
-            gp_model.eval()
-            likelihood.eval()
-            test_preds = likelihood(gp_model(test_x)).mean
-            mean_abs_error = torch.mean(torch.abs(test_y - test_preds))
-
-            self.assertLess(mean_abs_error.squeeze().item(), 0.02)
-
-            # Test variances
-            test_vars = likelihood(gp_model(test_x)).variance
-            self.assertAllClose(test_vars, likelihood(gp_model(test_x)).covariance_matrix.diagonal(dim1=-1, dim2=-2))
-            self.assertGreater(test_vars.min().item() + 0.1, likelihood.noise.item())
-            self.assertLess(
-                test_vars.max().item() - 0.05,
-                likelihood.noise.item() + gp_model.covar_module.base_kernel.outputscale.item()
-            )
+        with least_used_cuda_device():
+            self.test_sgpr_mean_abs_error(cuda=True)
 
 
 if __name__ == "__main__":

test/kernels/test_additive_and_product_kernels.py

@@ -56,6 +56,23 @@ def test_computes_product_of_radial_basis_function(self):
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 2e-5)
 
+    def test_computes_product_of_radial_basis_function_batch(self):
+        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
+        lengthscale = 2
+
+        kernel_1 = RBFKernel(batch_shape=torch.Size([4])).initialize(lengthscale=lengthscale)
+        kernel_2 = RBFKernel().initialize(lengthscale=lengthscale)
+        kernel = kernel_1 * kernel_2
+
+        actual = torch.tensor([[16, 4], [4, 0], [64, 36]], dtype=torch.float)
+        actual = actual.mul_(-0.5).div_(lengthscale ** 2).exp() ** 2
+        actual = actual.repeat(4, 1, 1)
+
+        kernel.eval()
+        res = kernel(a, b).evaluate()
+        self.assertLess(torch.norm(res - actual), 2e-5)
+
     def test_computes_sum_of_radial_basis_function(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
         b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)