Merge pull request #1674 from wjmaddox/sum_kronecker_iqldt

inv_quad_logdet method for sumKronecker

Author: Balandat

gpytorch/lazy/sum_kronecker_lazy_tensor.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 
-from .. import settings
 from .kronecker_product_lazy_tensor import KroneckerProductLazyTensor
 from .sum_lazy_tensor import SumLazyTensor
 
@@ -36,9 +35,6 @@ def _sum_formulation(self):
         return inv_root_times_lt1
 
     def _solve(self, rhs, preconditioner=None, num_tridiag=0):
-        if self.shape[-1] <= settings.max_cholesky_size.value():
-            return super()._solve(rhs=rhs, preconditioner=preconditioner, num_tridiag=num_tridiag)
-
         inner_mat = self._sum_formulation
         # root decomposition may not be trustworthy if it uses a different method than
         # root_inv_decomposition. so ensure that we call this locally
@@ -73,3 +69,19 @@ def _root_inv_decomposition(self, initial_vectors=None):
         inner_mat_root_inv = inner_mat.root_inv_decomposition().root
         inv_root = lt2_root_inv.matmul(inner_mat_root_inv)
         return inv_root
+
+    def inv_quad_logdet(self, inv_quad_rhs=None, logdet=False, reduce_inv_quad=True):
+        inv_quad_term = None
+        logdet_term = None
+
+        if inv_quad_rhs is not None:
+            solve = self.inv_matmul(inv_quad_rhs)
+            inv_quad_term = (inv_quad_rhs * solve).sum(-2)
+
+            if inv_quad_term.numel() and reduce_inv_quad:
+                inv_quad_term = inv_quad_term.sum(-1)
+
+        if logdet:
+            logdet_term = self._logdet()
+
+        return inv_quad_term, logdet_term

gpytorch/likelihoods/multitask_gaussian_likelihood.py

@@ -12,6 +12,7 @@
     DiagLazyTensor,
     KroneckerProductDiagLazyTensor,
     KroneckerProductLazyTensor,
+    LazyEvaluatedKernelTensor,
     RootLazyTensor,
 )
 from ..likelihoods import Likelihood, _GaussianLikelihoodBase
@@ -88,6 +89,10 @@ def marginal(self, function_dist, *params, **kwargs):
         """
         mean, covar = function_dist.mean, function_dist.lazy_covariance_matrix
 
+        # ensure that sumKroneckerLT is actually called
+        if isinstance(covar, LazyEvaluatedKernelTensor):
+            covar = covar.evaluate_kernel()
+
         covar_kron_lt = self._shaped_noise_covar(mean.shape, add_noise=self.has_global_noise)
         covar = covar + covar_kron_lt
 
@@ -126,7 +131,7 @@ def _shaped_noise_covar(self, shape, add_noise=True, *params, **kwargs):
 
     def forward(self, function_samples: Tensor, *params: Any, **kwargs: Any) -> base_distributions.Normal:
         noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()
-        noise = noise.view(*noise.shape[:-1], *function_samples.shape[-2:])
+        noise = noise.reshape(*noise.shape[:-1], *function_samples.shape[-2:])
         return base_distributions.Independent(base_distributions.Normal(function_samples, noise.sqrt()), 1)
 
 

test/lazy/test_sum_kronecker_lazy_tensor.py

@@ -22,7 +22,7 @@ class TestSumKroneckerLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_call_lanczos = True
     should_call_cg = False
-    skip_slq_tests = True
+    skip_slq_tests = False
 
     def create_lazy_tensor(self):
         a = torch.tensor([[4, 0, 2], [0, 3, -1], [2, -1, 3]], dtype=torch.float)
@@ -47,3 +47,19 @@ def evaluate_lazy_tensor(self, lazy_tensor):
             lazy_tensor.lazy_tensors[1].lazy_tensors[0].tensor, lazy_tensor.lazy_tensors[1].lazy_tensors[1].tensor
         )
         return res1 + res2
+
+    def test_inv_quad_logdet(self):
+        # mock call cg here
+        self.__class__.should_call_cg = True
+        super().test_inv_quad_logdet()
+        self.__class__.should_call_cg = False
+
+    def test_inv_quad_logdet_no_reduce(self):
+        self.__class__.should_call_cg = True
+        super().test_inv_quad_logdet_no_reduce()
+        self.__class__.should_call_cg = False
+
+    def test_root_decomposition_cholesky(self):
+        self.__class__.should_call_cg = True
+        super().test_root_decomposition_cholesky()
+        self.__class__.should_call_cg = False