inv_quad_logdet method for sumKronecker

Author: wjmaddox

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._solve(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

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