Repository: cornellius-gp/gpytorch
Pull request #614 (2019-04-03 21:38:14)

Root inv decomp cholesky
Fix to make `get_fantasy_model` work in small data regime

jacobrgardner: I'm like 66% sure the test failures here are actually PyTorch's fault, not ours: https://github.com/pytorch/pytorch/issues/18825
gpleiss: For now we could just add a backward hook that does `grad = 0.5 * (grad + grad.transpose(-1, -2)`, which will work even once the fix to pytorch is in place.

Patch:
diff --git a/gpytorch/functions/_inv_matmul.py b/gpytorch/functions/_inv_matmul.py
index 09bd8066..9077761c 100644
--- a/gpytorch/functions/_inv_matmul.py
+++ b/gpytorch/functions/_inv_matmul.py
@@ -90,7 +90,13 @@ class InvMatmul(Function):
                 left_solves = InvMatmul.apply(ctx.representation_tree, False, grad_output, *matrix_args)
 
                 if any(ctx.needs_input_grad[3:]):
-                    arg_grads = lazy_tsr._quad_form_derivative(left_solves, right_solves.mul(-1))
+                    # We call _quad_form_derivative to compute dl/dK
+                    # To ensure that this term is symmetric, we concatenate the left and right solves together,
+                    # and divide the result by 1/2
+                    arg_grads = lazy_tsr._quad_form_derivative(
+                        torch.cat([left_solves, right_solves], -1),
+                        torch.cat([right_solves, left_solves], -1).mul(-0.5)
+                    )
                 if ctx.needs_input_grad[2]:
                     right_grad = left_solves
                     if ctx.is_vector:
@@ -104,7 +110,11 @@ class InvMatmul(Function):
                 if ctx.needs_input_grad[3]:
                     left_grad = grad_output @ right_solves.transpose(-1, -2)
                 if any(ctx.needs_input_grad[4:]):
-                    arg_grads = lazy_tsr._quad_form_derivative(left_solves, right_solves.mul(-1))
+                    # We do this concatenation to ensure that the gradient of lazy_tsr is symmetric
+                    arg_grads = lazy_tsr._quad_form_derivative(
+                        torch.cat([left_solves, right_solves], -1),
+                        torch.cat([right_solves, left_solves], -1).mul(-0.5)
+                    )
                 if ctx.needs_input_grad[2]:
                     right_grad = left_solves
                     if ctx.is_vector:
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 629e72a2..5d8e12a6 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -17,6 +17,7 @@ from ..functions._root_decomposition import RootDecomposition
 from ..utils.broadcasting import _matmul_broadcast_shape, _mul_broadcast_shape
 from ..utils.cholesky import psd_safe_cholesky, cholesky_solve
 from ..utils.deprecation import _deprecate_renamed_methods
+from ..utils.gradients import _ensure_symmetric_grad
 from ..utils.getitem import _noop_index, _convert_indices_to_tensors, _compute_getitem_size
 from ..utils.memoize import add_to_cache, cached
 from ..utils.qr import batch_qr
@@ -392,7 +393,14 @@ class LazyTensor(ABC):
             (LazyTensor) Cholesky factor
         """
         from .non_lazy_tensor import NonLazyTensor
-        cholesky = psd_safe_cholesky(self.evaluate().double()).to(self.dtype)
+        evaluated_mat = self.evaluate()
+
+        # NOTE: this hack is in place so that the gradient of the Cholesky factorization is symmetric
+        # We can remove this hack once https://github.com/pytorch/pytorch/issues/18825 is merged in
+        if evaluated_mat.requires_grad:
+            evaluated_mat.register_hook(_ensure_symmetric_grad)
+
+        cholesky = psd_safe_cholesky(evaluated_mat.double()).to(self.dtype)
         return NonLazyTensor(cholesky)
 
     def _cholesky_solve(self, rhs):
@@ -697,6 +705,21 @@ class LazyTensor(ABC):
         """
         return self.shape[:-2]
 
+    def cholesky(self, upper=False):
+        """
+        Cholesky-factorizes the LazyTensor
+
+        Parameters:
+            upper (bool) - upper triangular or lower triangular factor (default: False)
+
+        Returns:
+            (LazyTensor) Cholesky factor (lower triangular)
+        """
+        res = self._cholesky()
+        if upper:
+            res = res.transpose(-1, -2)
+        return res
+
     def clone(self):
         """
         Clones the LazyTensor (creates clones of all underlying tensors)
@@ -966,7 +989,7 @@ class LazyTensor(ABC):
         if settings.fast_computations.log_prob.off() or (self.size(-1) <= settings.max_cholesky_size.value()):
             from .chol_lazy_tensor import CholLazyTensor
 
-            cholesky = CholLazyTensor(self._cholesky())
+            cholesky = CholLazyTensor(self.cholesky())
             return cholesky.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad)
 
         # Default: use modified batch conjugate gradients to compute these terms
@@ -1281,7 +1304,7 @@ class LazyTensor(ABC):
             or settings.fast_computations.covar_root_decomposition.off()
         ):
             try:
-                res = self._cholesky()
+                res = self.cholesky()
                 return CholLazyTensor(res)
 
             except RuntimeError as e:
@@ -1300,9 +1323,20 @@ class LazyTensor(ABC):
         low-rank version of a matrix
         """
         from .root_lazy_tensor import RootLazyTensor
+        from .non_lazy_tensor import lazify
 
-        if self.shape[-2:].numel() == 1:
-            return RootLazyTensor(1 / self.evaluate().sqrt())
+        if (
+            self.size(-1) <= settings.max_cholesky_size.value()
+            or settings.fast_computations.covar_root_decomposition.off()
+        ):
+            try:
+                res = self.cholesky()
+                res = lazify(delazify(res).double().inverse().transpose(-2, -1).to(self.dtype))
+                return RootLazyTensor(res)
+            except RuntimeError as e:
+                warnings.warn(
+                    "Runtime Error when computing Cholesky decomposition: {}. Using RootDecomposition.".format(e)
+                )
 
         if not self.is_square:
             raise RuntimeError(
diff --git a/gpytorch/utils/gradients.py b/gpytorch/utils/gradients.py
new file mode 100644
index 00000000..57e52555
--- /dev/null
+++ b/gpytorch/utils/gradients.py
@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+
+import torch
+
+
+def _ensure_symmetric_grad(grad):
+    """
+    A gradient-hook hack to ensure that symmetric matrix gradients are symmetric
+    """
+    res = torch.add(grad, grad.transpose(-1, -2)).mul(0.5)
+    return res
diff --git a/gpytorch/variational/variational_strategy.py b/gpytorch/variational/variational_strategy.py
index 8076978b..936537af 100644
--- a/gpytorch/variational/variational_strategy.py
+++ b/gpytorch/variational/variational_strategy.py
@@ -134,7 +134,7 @@ class VariationalStrategy(Module):
             # If we're less than a certain size, we'll compute the Cholesky decomposition of induc_induc_covar
             cholesky = False
             if settings.fast_computations.log_prob.off() or (num_induc <= settings.max_cholesky_size.value()):
-                induc_induc_covar = CholLazyTensor(induc_induc_covar._cholesky())
+                induc_induc_covar = CholLazyTensor(induc_induc_covar.cholesky())
                 cholesky = True
 
             # Cache the CG results
diff --git a/gpytorch/variational/whitened_variational_strategy.py b/gpytorch/variational/whitened_variational_strategy.py
index 8c137723..8c1f5a31 100644
--- a/gpytorch/variational/whitened_variational_strategy.py
+++ b/gpytorch/variational/whitened_variational_strategy.py
@@ -121,7 +121,7 @@ class WhitenedVariationalStrategy(VariationalStrategy):
             # If we're less than a certain size, we'll compute the Cholesky decomposition of induc_induc_covar
             cholesky = False
             if settings.fast_computations.log_prob.off() or (num_induc <= settings.max_cholesky_size.value()):
-                induc_induc_covar = CholLazyTensor(induc_induc_covar._cholesky())
+                induc_induc_covar = CholLazyTensor(induc_induc_covar.cholesky())
                 cholesky = True
 
             # Cache the CG results


Test patch:
diff --git a/test/functions/test_inv_matmul.py b/test/functions/test_inv_matmul.py
index 1d7460e7..bdf939e4 100644
--- a/test/functions/test_inv_matmul.py
+++ b/test/functions/test_inv_matmul.py
@@ -2,141 +2,70 @@
 
 import torch
 import unittest
-import os
-import random
 from gpytorch import settings
 from gpytorch.lazy import NonLazyTensor
+from gpytorch.utils.gradients import _ensure_symmetric_grad
+from test._base_test_case import BaseTestCase
 
 
-class TestInvMatmulNonBatch(unittest.TestCase):
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
-
-    def setUp(self):
-        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(1)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(1)
-            random.seed(1)
+class TestInvMatmulNonBatch(BaseTestCase, unittest.TestCase):
+    seed = 0
 
+    def _create_mat(self):
         mat = torch.randn(8, 8)
         mat = mat @ mat.transpose(-1, -2)
-        vec = torch.randn(8)
-        vecs = torch.randn(8, 4)
-        self.mat = mat.detach().clone().requires_grad_(True)
-        self.mat_copy = mat.detach().clone().requires_grad_(True)
-        self.vec = vec.detach().clone().requires_grad_(True)
-        self.vec_copy = vec.detach().clone().requires_grad_(True)
-        self.vecs = vecs.detach().clone().requires_grad_(True)
-        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
+        return mat
 
     def test_inv_matmul_vec(self):
-        # Forward
-        with settings.terminate_cg_by_size(False):
-            res = NonLazyTensor(self.mat).inv_matmul(self.vec)
-            actual = self.mat_copy.inverse().matmul(self.vec_copy)
-            self.assertLess(torch.max((res - actual).abs() / actual.abs()).item(), 1e-3)
-
-            # Backward
-            grad_output = torch.randn(8)
-            res.backward(gradient=grad_output)
-            actual.backward(gradient=grad_output)
-            self.assertLess(torch.max((self.mat_copy.grad - self.mat.grad).abs()).item(), 1e-3)
-            self.assertLess(torch.max((self.vec_copy.grad - self.vec.grad).abs()).item(), 1e-3)
+        mat = self._create_mat().detach().requires_grad_(True)
+        if mat.dim() > 2:  # This isn't a feature for batch mode
+            return
+        mat_copy = mat.detach().clone().requires_grad_(True)
+        mat_copy.register_hook(_ensure_symmetric_grad)
+        vec = torch.randn(mat.size(-1)).detach().requires_grad_(True)
+        vec_copy = vec.detach().clone().requires_grad_(True)
 
-    def test_inv_matmul_multiple_vecs(self):
         # Forward
         with settings.terminate_cg_by_size(False):
-            res = NonLazyTensor(self.mat).inv_matmul(self.vecs)
-            actual = self.mat_copy.inverse().matmul(self.vecs_copy)
-            self.assertLess(torch.max((res - actual).abs() / actual.abs()).item(), 1e-3)
+            res = NonLazyTensor(mat).inv_matmul(vec)
+            actual = mat_copy.inverse().matmul(vec_copy)
+            self.assertAllClose(res, actual)
 
             # Backward
-            grad_output = torch.randn(8, 4)
+            grad_output = torch.randn_like(vec)
             res.backward(gradient=grad_output)
             actual.backward(gradient=grad_output)
-            self.assertLess(torch.max((self.mat_copy.grad - self.mat.grad).abs()).item(), 1e-3)
-            self.assertLess(torch.max((self.vecs_copy.grad - self.vecs.grad).abs()).item(), 1e-3)
-
-
-class TestInvMatmulBatch(unittest.TestCase):
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
-
-    def setUp(self):
-        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(0)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(0)
-            random.seed(0)
-
-        mats = torch.randn(2, 8, 8)
-        mats = mats @ mats.transpose(-1, -2)
-        vecs = torch.randn(2, 8, 4)
-        self.mats = mats.detach().clone().requires_grad_(True)
-        self.mats_copy = mats.detach().clone().requires_grad_(True)
-        self.vecs = vecs.detach().clone().requires_grad_(True)
-        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
+            self.assertAllClose(mat.grad, mat_copy.grad)
+            self.assertAllClose(vec.grad, vec_copy.grad)
 
     def test_inv_matmul_multiple_vecs(self):
+        mat = self._create_mat().detach().requires_grad_(True)
+        mat_copy = mat.detach().clone().requires_grad_(True)
+        mat_copy.register_hook(_ensure_symmetric_grad)
+        vecs = torch.randn(*mat.shape[:-2], mat.size(-1), 4).detach().requires_grad_(True)
+        vecs_copy = vecs.detach().clone().requires_grad_(True)
+
         # Forward
         with settings.terminate_cg_by_size(False):
-            res = NonLazyTensor(self.mats).inv_matmul(self.vecs)
-            actual = torch.cat(
-                [self.mats_copy[0].inverse().unsqueeze(0), self.mats_copy[1].inverse().unsqueeze(0)]
-            ).matmul(self.vecs_copy)
-            self.assertLess(torch.max((res - actual).abs()).item(), 1e-3)
+            res = NonLazyTensor(mat).inv_matmul(vecs)
+            actual = mat_copy.inverse().matmul(vecs_copy)
+            self.assertAllClose(res, actual)
 
             # Backward
-            grad_output = torch.randn(2, 8, 4)
+            grad_output = torch.randn_like(vecs)
             res.backward(gradient=grad_output)
             actual.backward(gradient=grad_output)
-            self.assertLess(torch.max((self.mats_copy.grad - self.mats.grad).abs()).item(), 1e-3)
-            self.assertLess(torch.max((self.vecs_copy.grad - self.vecs.grad).abs()).item(), 1e-3)
-
+            self.assertAllClose(mat.grad, mat_copy.grad)
+            self.assertAllClose(vecs.grad, vecs_copy.grad)
 
-class TestInvMatmulMultiBatch(unittest.TestCase):
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
 
-    def setUp(self):
-        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(0)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(0)
-            random.seed(0)
+class TestInvMatmulBatch(TestInvMatmulNonBatch):
+    seed = 0
 
-        mats = torch.randn(3, 4, 8, 8)
+    def _create_mat(self):
+        mats = torch.randn(2, 8, 8)
         mats = mats @ mats.transpose(-1, -2)
-        mats.add_(torch.eye(8).mul_(1e-1).view(1, 1, 8, 8))
-        vecs = torch.randn(3, 4, 8, 2)
-        self.mats = mats.detach().clone().requires_grad_(True)
-        self.mats_copy = mats.detach().clone().requires_grad_(True)
-        self.vecs = vecs.detach().clone().requires_grad_(True)
-        self.vecs_copy = vecs.detach().clone().requires_grad_(True)
-
-    def test_inv_matmul_multiple_vecs(self):
-        # Forward
-        with settings.terminate_cg_by_size(False):
-            res = NonLazyTensor(self.mats).inv_matmul(self.vecs)
-            flattened_mats_copy = self.mats_copy.view(-1, *self.mats.shape[-2:])
-            flatened_mats_inverse = torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats_copy])
-            mats_inverse = flatened_mats_inverse.view_as(self.mats)
-            actual = mats_inverse.matmul(self.vecs_copy)
-            self.assertLess(torch.max((res - actual).abs()).item(), 1e-3)
-
-            # Backward
-            grad_output = torch.randn(3, 4, 8, 2)
-            res.backward(gradient=grad_output)
-            actual.backward(gradient=grad_output)
-            self.assertLess(torch.max((self.mats_copy.grad - self.mats.grad).abs()).item(), 1e-3)
-            self.assertLess(torch.max((self.vecs_copy.grad - self.vecs.grad).abs()).item(), 1e-3)
+        return mats
 
 
 if __name__ == "__main__":
diff --git a/test/functions/test_root_decomposition.py b/test/functions/test_root_decomposition.py
index 8bb692eb..097c974e 100644
--- a/test/functions/test_root_decomposition.py
+++ b/test/functions/test_root_decomposition.py
@@ -1,148 +1,71 @@
 #!/usr/bin/env python3
 
-import os
-import random
 import unittest
-from test._utils import approx_equal
+from .._base_test_case import BaseTestCase
 
 import torch
 from gpytorch.lazy import NonLazyTensor
 
 
-class TestRootDecomposition(unittest.TestCase):
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
-
-    def setUp(self):
-        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(0)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(0)
-            random.seed(0)
+class TestRootDecomposition(BaseTestCase, unittest.TestCase):
+    seed = 0
 
+    def _create_mat(self):
         mat = torch.randn(4, 4)
         mat = mat @ mat.transpose(-1, -2)
         mat.div_(5).add_(torch.eye(4))
-        self.mat = mat.detach().clone().requires_grad_(True)
-        self.mat_clone = mat.detach().clone().requires_grad_(True)
+        return mat
 
     def test_root_decomposition(self):
+        mat = self._create_mat().detach().requires_grad_(True)
+        mat_clone = mat.detach().clone().requires_grad_(True)
+
         # Forward
-        root = NonLazyTensor(self.mat).root_decomposition().root.evaluate()
+        root = NonLazyTensor(mat).root_decomposition().root.evaluate()
         res = root.matmul(root.transpose(-1, -2))
-        self.assertTrue(approx_equal(res, self.mat))
+        self.assertAllClose(res, mat)
 
         # Backward
-        res.trace().backward()
-        self.mat_clone.trace().backward()
-        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
+        sum([mat.trace() for mat in res.view(-1, mat.size(-2), mat.size(-1))]).backward()
+        sum([mat.trace() for mat in mat_clone.view(-1, mat.size(-2), mat.size(-1))]).backward()
+        self.assertAllClose(mat.grad, mat_clone.grad)
 
     def test_root_inv_decomposition(self):
+        mat = self._create_mat().detach().requires_grad_(True)
+        mat_clone = mat.detach().clone().requires_grad_(True)
+
         # Forward
-        probe_vectors = torch.randn(4, 5)
-        test_vectors = torch.randn(4, 5)
-        root = NonLazyTensor(self.mat).root_inv_decomposition(probe_vectors, test_vectors).root.evaluate()
+        probe_vectors = torch.randn(*mat.shape[:-2], 4, 5)
+        test_vectors = torch.randn(*mat.shape[:-2], 4, 5)
+        root = NonLazyTensor(mat).root_inv_decomposition(probe_vectors, test_vectors).root.evaluate()
         res = root.matmul(root.transpose(-1, -2))
-        actual = self.mat_clone.inverse()
-        self.assertTrue(approx_equal(res, actual))
+        actual = mat_clone.inverse()
+        self.assertAllClose(res, actual)
 
         # Backward
-        res.trace().backward()
-        actual.trace().backward()
-        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
-
+        sum([mat.trace() for mat in res.view(-1, mat.size(-2), mat.size(-1))]).backward()
+        sum([mat.trace() for mat in actual.view(-1, mat.size(-2), mat.size(-1))]).backward()
+        self.assertAllClose(mat.grad, mat_clone.grad)
 
-class TestRootDecompositionBatch(unittest.TestCase):
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
 
-    def setUp(self):
-        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(0)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(0)
-            random.seed(0)
+class TestRootDecompositionBatch(TestRootDecomposition):
+    seed = 0
 
+    def _create_mat(self):
         mat = torch.randn(3, 4, 4)
         mat = mat @ mat.transpose(-1, -2)
         mat.div_(5).add_(torch.eye(4).unsqueeze_(0))
-        self.mat = mat.detach().clone().requires_grad_(True)
-        self.mat_clone = mat.detach().clone().requires_grad_(True)
-
-    def test_root_decomposition(self):
-        # Forward
-        root = NonLazyTensor(self.mat).root_decomposition().root.evaluate()
-        res = root.matmul(root.transpose(-1, -2))
-        self.assertTrue(approx_equal(res, self.mat))
-
-        # Backward
-        sum([mat.trace() for mat in res]).backward()
-        sum([mat.trace() for mat in self.mat_clone]).backward()
-        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
-
-    def test_root_inv_decomposition(self):
-        # Forward
-        probe_vectors = torch.randn(3, 4, 5)
-        test_vectors = torch.randn(3, 4, 5)
-        root = NonLazyTensor(self.mat).root_inv_decomposition(probe_vectors, test_vectors).root.evaluate()
-        res = root.matmul(root.transpose(-1, -2))
-        actual = torch.cat([mat.inverse().unsqueeze(0) for mat in self.mat_clone])
-        self.assertTrue(approx_equal(res, actual))
-
-        # Backward
-        sum([mat.trace() for mat in res]).backward()
-        sum([mat.trace() for mat in actual]).backward()
-        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
-
+        return mat
 
-class TestRootDecompositionMultiBatch(unittest.TestCase):
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
 
-    def setUp(self):
-        if os.getenv("unlock_seed") is None or os.getenv("unlock_seed").lower() == "false":
-            self.rng_state = torch.get_rng_state()
-            torch.manual_seed(0)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(0)
-            random.seed(0)
+class TestRootDecompositionMultiBatch(TestRootDecomposition):
+    seed = 0
 
+    def _create_mat(self):
         mat = torch.randn(2, 3, 4, 4)
         mat = mat @ mat.transpose(-1, -2)
-        mat.div_(5).add_(torch.eye(4).view(1, 1, 4, 4))
-        self.mat = mat.detach().clone().requires_grad_(True)
-        self.mat_clone = mat.detach().clone().requires_grad_(True)
-
-    def test_root_decomposition(self):
-        # Forward
-        root = NonLazyTensor(self.mat).root_decomposition().root.evaluate()
-        res = root.matmul(root.transpose(-1, -2))
-        self.assertTrue(approx_equal(res, self.mat))
-
-        # Backward
-        sum([mat.trace() for mat in res.view(-1, *self.mat.shape[-2:])]).backward()
-        sum([mat.trace() for mat in self.mat_clone.view(-1, *self.mat.shape[-2:])]).backward()
-        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
-
-    def test_root_inv_decomposition(self):
-        # Forward
-        probe_vectors = torch.randn(2, 3, 4, 5)
-        test_vectors = torch.randn(2, 3, 4, 5)
-        root = NonLazyTensor(self.mat).root_inv_decomposition(probe_vectors, test_vectors).root.evaluate()
-        res = root.matmul(root.transpose(-1, -2))
-        flattened_mats = self.mat_clone.view(-1, *self.mat_clone.shape[-2:])
-        actual = torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats]).view_as(self.mat_clone)
-        self.assertTrue(approx_equal(res, actual))
-
-        # Backward
-        sum([mat.trace() for mat in res.view(-1, *self.mat.shape[-2:])]).backward()
-        sum([mat.trace() for mat in actual.view(-1, *self.mat.shape[-2:])]).backward()
-        self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
+        mat.div_(5).add_(torch.eye(4).unsqueeze_(0))
+        return mat
 
 
 if __name__ == "__main__":
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index a494286a..18fe9f14 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch
 
 import gpytorch
 import torch
+from gpytorch.utils.gradients import _ensure_symmetric_grad
 from test._base_test_case import BaseTestCase
 
 
@@ -248,6 +249,7 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
         lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
         lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
         evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+        evaluated.register_hook(_ensure_symmetric_grad)
 
         # Create a test right hand side and left hand side
         rhs.requires_grad_(True)
diff --git a/test/lazy/test_cached_cg_lazy_tensor.py b/test/lazy/test_cached_cg_lazy_tensor.py
index 45d7899c..d9cfe599 100644
--- a/test/lazy/test_cached_cg_lazy_tensor.py
+++ b/test/lazy/test_cached_cg_lazy_tensor.py
@@ -6,6 +6,7 @@ import gpytorch
 import unittest
 import warnings
 from gpytorch.lazy import CachedCGLazyTensor, NonLazyTensor
+from gpytorch.utils.gradients import _ensure_symmetric_grad
 from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 from unittest.mock import MagicMock, patch
 
@@ -47,6 +48,7 @@ class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
         lazy_tensor = self.create_lazy_tensor(with_solves=True).requires_grad_(True)
         lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
         evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
+        evaluated.register_hook(_ensure_symmetric_grad)
 
         # Rather than the supplied rhs and lhs,
         # we'll replace them with ones that we've precomputed solves for