Merge pull request #1394 from wjmaddox/kplt_general_roots

Make KroneckerProductLazyTensor call each component's root decomposition

Author: gpleiss

gpytorch/lazy/kronecker_product_lazy_tensor.py

@@ -199,18 +199,26 @@ def _matmul(self, rhs):
     def root_decomposition(self, method: Optional[str] = None):
         from gpytorch.lazy import RootLazyTensor
 
-        if method == "symeig" or method is None:
-            evals, evecs = self._symeig(eigenvectors=True, return_evals_as_lazy=True)
-            # TODO: only use non-zero evals (req. dealing w/ batches...)
-            f_list = [
-                evec * eval.diag().clamp(0.0).sqrt().unsqueeze(-2)
-                for eval, evec in zip(evals.lazy_tensors, evecs.lazy_tensors)
-            ]
-            F = KroneckerProductLazyTensor(*f_list)
-            return RootLazyTensor(F)
-        else:
+        # return a dense root decomposition if the matrix is small
+        if self.shape[-1] <= settings.max_cholesky_size.value():
             return super().root_decomposition(method=method)
 
+        root_list = [lt.root_decomposition(method=method).root for lt in self.lazy_tensors]
+        kronecker_root = KroneckerProductLazyTensor(*root_list)
+        return RootLazyTensor(kronecker_root)
+
+    @cached(name="root_inv_decomposition")
+    def root_inv_decomposition(self, initial_vectors=None, test_vectors=None):
+        from gpytorch.lazy import RootLazyTensor
+
+        # return a dense root decomposition if the matrix is small
+        if self.shape[-1] <= settings.max_cholesky_size.value():
+            return super().root_inv_decomposition()
+
+        root_list = [lt.root_inv_decomposition().root for lt in self.lazy_tensors]
+        kronecker_root = KroneckerProductLazyTensor(*root_list)
+        return RootLazyTensor(kronecker_root)
+
     @cached(name="size")
     def _size(self):
         left_size = _prod(lazy_tensor.size(-2) for lazy_tensor in self.lazy_tensors)

test/lazy/test_kronecker_product_lazy_tensor.py

@@ -20,7 +20,7 @@ def kron(a, b):
 
 class TestKroneckerProductLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
-    should_call_lanczos = False
+    should_call_lanczos = True
 
     def create_lazy_tensor(self):
         a = torch.tensor([[4, 0, 2], [0, 3, -1], [2, -1, 3]], dtype=torch.float)
@@ -40,12 +40,14 @@ def evaluate_lazy_tensor(self, lazy_tensor):
 
 class TestKroneckerProductLazyTensorBatch(TestKroneckerProductLazyTensor):
     seed = 0
-    should_call_lanczos = False
+    should_call_lanczos = True
 
     def create_lazy_tensor(self):
         a = torch.tensor([[4, 0, 2], [0, 3, -1], [2, -1, 3]], dtype=torch.float).repeat(3, 1, 1)
         b = torch.tensor([[2, 1], [1, 2]], dtype=torch.float).repeat(3, 1, 1)
-        c = torch.tensor([[4, 0, 1, 0], [0, 4, -1, 0], [1, -1, 3, 0], [0, 0, 0, 4]], dtype=torch.float).repeat(3, 1, 1)
+        c = torch.tensor([[4, 0.1, 1, 0], [0.1, 4, -1, 0], [1, -1, 3, 0], [0, 0, 0, 4]], dtype=torch.float).repeat(
+            3, 1, 1
+        )
         a.requires_grad_(True)
         b.requires_grad_(True)
         c.requires_grad_(True)