Repository: cornellius-gp/gpytorch
Pull request #586 (2019-03-19 21:00:05)

Better Cholesky methods for use in small data regimes and for debugging
This PR does a few things:

- [x] Introduces the `_cholesky`, `_cholesky_solve`, `_root_decomposition`, and `_root_inv_decomposition` private methods for LazyTensors. This makes it easy to define efficient Cholesky decompositions without having to re-write all the logic of `
- [x] Use Cholesky for `inv_matmul` (on by default for small matrices)
- [x] Use Cholesky for `inv_quad_logdet` (on by default for small matrices)
- [x] Make sure cholesky option is possible for VariationalStrategies

Cholesky is faster and more efficient than CG with `n <= 300` (see #585)


Patch:
diff --git a/.travis.yml b/.travis.yml
index df6d4559..160ac6ff 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -32,6 +32,6 @@ matrix:
       addons:
         apt_packages:
           - pandoc
-      python: "3.6"
+      python: "3.5"
       install: pip install IPython ipykernel sphinx sphinx_rtd_theme nbsphinx m2r
       script: python setup.py build_sphinx
diff --git a/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb b/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
index 4a3959bf..6a516047 100644
--- a/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
+++ b/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
@@ -6,12 +6,12 @@
    "source": [
     "# Exact GP Regression with Multiple GPUs and Kernel Partitioning\n",
     "\n",
-    "In this notebook, we'll demonstrate training exact GPs on large datasets using two key features from the paper https://arxiv.org/abs/1903.08114: \n",
+    "In this notebook, we'll demonstrate training exact GPs on very large datapoints using two key features:\n",
     "\n",
-    "1. The ability to distribute the kernel matrix across multiple GPUs, for additional parallelism.\n",
-    "2. Partitioning the kernel into chunks computed on-the-fly when performing each MVM to reduce memory usage.\n",
+    "1. The ability to distribute the kernel matrix across multiple GPUs, thereby using additional parallelism.\n",
+    "2. Partitioning the kernel in to chunks that are computed on-the-fly when performing each MVM so as to reduce memory usage enough.\n",
     "\n",
-    "We'll be using the `protein` dataset, which has about 37000 training examples. The techniques in this notebook can be applied to much larger datasets, but the training time required will depend on the computational resources you have available: both the number of GPUs available and the amount of memory they have (which determines the partition size) have a significant effect on training time."
+    "We'll be using the `protein` dataset, which will have about 37000 training examples. Note that although the techniques in this notebook can be applied to much larger datasets than this, the training time required will depend dramatically on the computational resources you have available: both the number of GPUs available and the amount of memory they have (which determines the partition size) have a dramatic effect on performance."
    ]
   },
   {
diff --git a/gpytorch/functions/_inv_matmul.py b/gpytorch/functions/_inv_matmul.py
index 959f9a9d..b6590581 100644
--- a/gpytorch/functions/_inv_matmul.py
+++ b/gpytorch/functions/_inv_matmul.py
@@ -5,6 +5,13 @@ from torch.autograd import Function
 from .. import settings
 
 
+def _solve(lazy_tsr, rhs, preconditioner):
+    if settings.fast_computations.solves.off() or lazy_tsr.size(-1) <= settings.max_cholesky_size.value():
+        return lazy_tsr._cholesky()._cholesky_solve(rhs)
+    else:
+        return lazy_tsr._solve(rhs, preconditioner)
+
+
 class InvMatmul(Function):
     def __init__(self, representation_tree, has_left=False):
         self.representation_tree = representation_tree
@@ -32,11 +39,11 @@ class InvMatmul(Function):
         # Perform solves (for inv_quad) and tridiagonalization (for estimating logdet)
         if self.has_left:
             rhs = torch.cat([left_tensor.transpose(-1, -2), right_tensor], -1)
-            solves = lazy_tsr._solve(rhs, self.preconditioner)
+            solves = _solve(lazy_tsr, rhs, self.preconditioner)
             res = solves[..., left_tensor.size(-2):]
             res = left_tensor @ res
         else:
-            solves = lazy_tsr._solve(right_tensor, self.preconditioner)
+            solves = _solve(lazy_tsr, right_tensor, self.preconditioner)
             res = solves
 
         if self.is_vector:
@@ -79,7 +86,7 @@ class InvMatmul(Function):
 
             if not self.has_left:
                 # Compute self^{-1} grad_output
-                left_solves = lazy_tsr._solve(grad_output, self.preconditioner)
+                left_solves = _solve(lazy_tsr, grad_output, self.preconditioner)
 
                 if any(self.needs_input_grad[1:]):
                     arg_grads = lazy_tsr._quad_form_derivative(left_solves, right_solves.mul(-1))
diff --git a/gpytorch/functions/_inv_quad.py b/gpytorch/functions/_inv_quad.py
new file mode 100644
index 00000000..53940314
--- /dev/null
+++ b/gpytorch/functions/_inv_quad.py
@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+
+import torch
+from torch.autograd import Function
+from .. import settings
+
+
+def _solve(lazy_tsr, rhs, preconditioner):
+    if settings.fast_computations.solves.off() or settings.fast_computations.log_prob.off() or \
+            lazy_tsr.size(-1) <= settings.max_cholesky_size.value():
+        return lazy_tsr._cholesky()._cholesky_solve(rhs)
+    else:
+        return lazy_tsr._solve(rhs, preconditioner)
+
+
+class InvQuad(Function):
+    """
+    Given a PSD matrix A (or a batch of PSD matrices A), this function computes b A^{-1} b
+    where b is a vector or batch of vectors
+    """
+
+    def __init__(self, representation_tree):
+        self.representation_tree = representation_tree
+
+    def forward(self, *args):
+        """
+        *args - The arguments representing the PSD matrix A (or batch of PSD matrices A)
+        If self.inv_quad is true, the first entry in *args is inv_quad_rhs (Tensor)
+        - the RHS of the matrix solves.
+
+        Returns:
+        - (Scalar) The inverse quadratic form (or None, if self.inv_quad is False)
+        - (Scalar) The log determinant (or None, self.if logdet is False)
+        """
+        inv_quad_rhs, *matrix_args = args
+
+        # Get closure for matmul
+        lazy_tsr = self.representation_tree(*matrix_args)
+        with torch.no_grad():
+            preconditioner = lazy_tsr.detach()._inv_matmul_preconditioner()
+
+        # RHS for inv_quad
+        self.is_vector = False
+        if inv_quad_rhs.ndimension() == 1:
+            inv_quad_rhs = inv_quad_rhs.unsqueeze(-1)
+            self.is_vector = True
+
+        # Perform solves (for inv_quad) and tridiagonalization (for estimating logdet)
+        inv_quad_solves = _solve(lazy_tsr, inv_quad_rhs, preconditioner)
+        inv_quad_term = (inv_quad_solves * inv_quad_rhs).sum(-2)
+
+        to_save = matrix_args + [inv_quad_solves]
+        self.save_for_backward(*to_save)
+
+        if settings.memory_efficient.off():
+            self._lazy_tsr = lazy_tsr
+
+        return inv_quad_term
+
+    def backward(self, inv_quad_grad_output):
+        *matrix_args, inv_quad_solves = self.saved_tensors
+
+        if hasattr(self, "_lazy_tsr"):
+            lazy_tsr = self._lazy_tsr
+        else:
+            lazy_tsr = self.representation_tree(*matrix_args)
+
+        # Fix grad_output sizes
+        inv_quad_grad_output = inv_quad_grad_output.unsqueeze(-2)
+        neg_inv_quad_solves_times_grad_out = inv_quad_solves.mul(inv_quad_grad_output).mul_(-1)
+
+        # input_1 gradient
+        if any(self.needs_input_grad[1:]):
+            left_factors = neg_inv_quad_solves_times_grad_out
+            right_factors = inv_quad_solves
+            matrix_arg_grads = lazy_tsr._quad_form_derivative(left_factors, right_factors)
+
+        # input_2 gradients
+        if self.needs_input_grad[0]:
+            inv_quad_rhs_grad = neg_inv_quad_solves_times_grad_out.mul_(-2)
+        else:
+            inv_quad_rhs_grad = torch.zeros_like(inv_quad_solves)
+        if self.is_vector:
+            inv_quad_rhs_grad.squeeze_(-1)
+
+        res = tuple([inv_quad_rhs_grad] + list(matrix_arg_grads))
+        return tuple(res)
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
index 336bd69e..d6cf5ab2 100644
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -51,6 +51,9 @@ class InvQuadLogDet(Function):
         self.probe_vectors = probe_vectors
         self.probe_vector_norms = probe_vector_norms
 
+        if self.logdet and not self.probe_vectors.numel():
+            raise RuntimeError("Probe vectors were not supplied for logdet computation")
+
     def forward(self, *args):
         """
         *args - The arguments representing the PSD matrix A (or batch of PSD matrices A)
diff --git a/gpytorch/functions/_root_decomposition.py b/gpytorch/functions/_root_decomposition.py
index a6df6781..7307f4a5 100644
--- a/gpytorch/functions/_root_decomposition.py
+++ b/gpytorch/functions/_root_decomposition.py
@@ -2,7 +2,7 @@
 
 import torch
 from torch.autograd import Function
-from ..utils.lanczos import lanczos_tridiag, lanczos_tridiag_to_diag
+from ..utils import lanczos
 from .. import settings
 
 
@@ -43,7 +43,7 @@ class RootDecomposition(Function):
         lazy_tsr = self.representation_tree(*matrix_args)
         matmul_closure = lazy_tsr._matmul
         # Do lanczos
-        q_mat, t_mat = lanczos_tridiag(
+        q_mat, t_mat = lanczos.lanczos_tridiag(
             matmul_closure,
             self.max_iter,
             dtype=self.dtype,
@@ -65,7 +65,7 @@ class RootDecomposition(Function):
         jitter_mat = (settings.tridiagonal_jitter.value() * mins) * torch.eye(
             t_mat.size(-1), device=t_mat.device, dtype=t_mat.dtype
         ).expand_as(t_mat)
-        eigenvalues, eigenvectors = lanczos_tridiag_to_diag(t_mat + jitter_mat)
+        eigenvalues, eigenvectors = lanczos.lanczos_tridiag_to_diag(t_mat + jitter_mat)
 
         # Get orthogonal matrix and eigenvalue roots
         q_mat = q_mat.matmul(eigenvectors)
diff --git a/gpytorch/lazy/batch_repeat_lazy_tensor.py b/gpytorch/lazy/batch_repeat_lazy_tensor.py
index 25ea3bc6..c6a47ba2 100644
--- a/gpytorch/lazy/batch_repeat_lazy_tensor.py
+++ b/gpytorch/lazy/batch_repeat_lazy_tensor.py
@@ -3,11 +3,9 @@
 import itertools
 import torch
 from .. import settings
-from ..utils.memoize import cached
 from ..utils.broadcasting import _matmul_broadcast_shape
-
+from ..utils.memoize import cached
 from .lazy_tensor import LazyTensor
-from .root_lazy_tensor import RootLazyTensor
 
 
 class BatchRepeatLazyTensor(LazyTensor):
@@ -34,6 +32,20 @@ class BatchRepeatLazyTensor(LazyTensor):
         self.base_lazy_tensor = base_lazy_tensor
         self.batch_repeat = batch_repeat
 
+    @cached(name="cholesky")
+    def _cholesky(self):
+        return self.base_lazy_tensor._cholesky().repeat(*self.batch_repeat, 1, 1)
+
+    def _cholesky_solve(self, rhs):
+        output_shape = _matmul_broadcast_shape(self.shape, rhs.shape)
+        if rhs.shape != output_shape:
+            rhs = rhs.expand(*output_shape)
+
+        rhs = self._move_repeat_batches_to_columns(rhs, output_shape)
+        res = self.base_lazy_tensor._cholesky_solve(rhs)
+        res = self._move_repeat_batches_back(res, output_shape)
+        return res
+
     def _compute_batch_repeat_size(self, current_batch_shape, desired_batch_shape):
         batch_repeat = torch.Size(
             desired_batch_size // current_batch_size
@@ -169,6 +181,12 @@ class BatchRepeatLazyTensor(LazyTensor):
         right_vectors = self._move_repeat_batches_to_columns(right_vectors, right_output_shape)
         return self.base_lazy_tensor._quad_form_derivative(left_vectors, right_vectors)
 
+    def _root_decomposition(self):
+        return self.base_lazy_tensor._root_decomposition().repeat(*self.batch_repeat, 1, 1)
+
+    def _root_inv_decomposition(self, initial_vectors=None):
+        return self.base_lazy_tensor._root_inv_decomposition().repeat(*self.batch_repeat, 1, 1)
+
     def _size(self):
         repeated_batch_shape = torch.Size(
             size * repeat for size, repeat in zip(self.base_lazy_tensor.batch_shape, self.batch_repeat)
@@ -250,14 +268,3 @@ class BatchRepeatLazyTensor(LazyTensor):
                 for orig_repeat_size, new_repeat_size in zip(padded_batch_repeat, sizes[:-2])
             ),
         )
-
-    def root_decomposition(self):
-        return RootLazyTensor(self.base_lazy_tensor.root_decomposition().root.repeat(*self.batch_repeat, 1, 1))
-
-    @cached
-    def root_inv_decomposition(self, initial_vectors=None, test_vectors=None):
-        return RootLazyTensor(
-            self.base_lazy_tensor.root_inv_decomposition(
-                initial_vectors=initial_vectors, test_vectors=test_vectors
-            ).root.repeat(*self.batch_repeat, 1, 1)
-        )
diff --git a/gpytorch/lazy/block_diag_lazy_tensor.py b/gpytorch/lazy/block_diag_lazy_tensor.py
index 932dd76e..9de9bab4 100644
--- a/gpytorch/lazy/block_diag_lazy_tensor.py
+++ b/gpytorch/lazy/block_diag_lazy_tensor.py
@@ -1,13 +1,8 @@
 #!/usr/bin/env python3
 
 import torch
-
-from .. import settings
-from ..utils.cholesky import psd_safe_cholesky
-from ..utils.memoize import cached
 from .block_lazy_tensor import BlockLazyTensor
-from .non_lazy_tensor import NonLazyTensor
-from .root_lazy_tensor import RootLazyTensor
+from ..utils.memoize import cached
 
 
 class BlockDiagLazyTensor(BlockLazyTensor):
@@ -36,6 +31,16 @@ class BlockDiagLazyTensor(BlockLazyTensor):
         other = other.view(*batch_shape, num_rows // self.num_blocks, num_cols)
         return other
 
+    @cached(name="cholesky")
+    def _cholesky(self):
+        return self.__class__(self.base_lazy_tensor._cholesky())
+
+    def _cholesky_solve(self, rhs):
+        rhs = self._add_batch_dim(rhs)
+        res = self.base_lazy_tensor._cholesky_solve(rhs)
+        res = self._remove_batch_dim(res)
+        return res
+
     def _get_indices(self, row_index, col_index, *batch_indices):
         # Figure out what block the row/column indices belong to
         row_index_block = row_index.div(self.base_lazy_tensor.size(-2))
@@ -58,6 +63,12 @@ class BlockDiagLazyTensor(BlockLazyTensor):
         other = other.contiguous().view(*shape)
         return other
 
+    def _root_decomposition(self):
+        return self.__class__(self.base_lazy_tensor._root_decomposition())
+
+    def _root_inv_decomposition(self, initial_vectors=None):
+        return self.__class__(self.base_lazy_tensor._root_inv_decomposition(initial_vectors))
+
     def _size(self):
         shape = list(self.base_lazy_tensor.shape)
         shape[-2] *= shape[-3]
@@ -65,6 +76,15 @@ class BlockDiagLazyTensor(BlockLazyTensor):
         del shape[-3]
         return torch.Size(shape)
 
+    def _solve(self, rhs, preconditioner, num_tridiag=0):
+        if num_tridiag:
+            return super()._solve(rhs, preconditioner, num_tridiag=num_tridiag)
+        else:
+            rhs = self._add_batch_dim(rhs)
+            res = self.base_lazy_tensor._solve(rhs, preconditioner, num_tridiag=None)
+            res = self._remove_batch_dim(res)
+            return res
+
     def diag(self):
         res = self.base_lazy_tensor.diag().contiguous()
         return res.view(*self.batch_shape, self.size(-1))
@@ -85,12 +105,3 @@ class BlockDiagLazyTensor(BlockLazyTensor):
         if logdet_res is not None and logdet_res.numel():
             logdet_res = logdet_res.view(*logdet_res.shape).sum(-1)
         return inv_quad_res, logdet_res
-
-    @cached(name="root_decomposition")
-    def root_decomposition(self):
-        if settings.fast_computations.covar_root_decomposition.on():
-            res = self.__class__(self.base_lazy_tensor.root_decomposition().root)
-        else:
-            chol = psd_safe_cholesky(self.base_lazy_tensor.evaluate())
-            res = self.__class__(NonLazyTensor(chol))
-        return RootLazyTensor(res)
diff --git a/gpytorch/lazy/cached_cg_lazy_tensor.py b/gpytorch/lazy/cached_cg_lazy_tensor.py
index 79eda57a..7ea0f0e3 100644
--- a/gpytorch/lazy/cached_cg_lazy_tensor.py
+++ b/gpytorch/lazy/cached_cg_lazy_tensor.py
@@ -75,7 +75,10 @@ class CachedCGLazyTensor(LazyTensor):
 
             else:
                 # Compute solves
-                solves = base_lazy_tensor._solve(eager_rhs, preconditioner=base_lazy_tensor._preconditioner()[0])
+                if settings.fast_computations.log_prob.on():
+                    solves = base_lazy_tensor._solve(eager_rhs, preconditioner=base_lazy_tensor._preconditioner()[0])
+                else:
+                    solves = base_lazy_tensor._cholesky()._cholesky_solve(eager_rhs)
                 dtype = solves.dtype
                 device = solves.device
                 return (
@@ -109,6 +112,9 @@ class CachedCGLazyTensor(LazyTensor):
     def requires_grad(self, val):
         self.base_lazy_tensor.requires_grad = val
 
+    def _cholesky(self):
+        return self.base_lazy_tensor._cholesky()
+
     def _expand_batch(self, batch_shape):
         return self.base_lazy_tensor._expand_batch(batch_shape)
 
@@ -127,7 +133,7 @@ class CachedCGLazyTensor(LazyTensor):
     def _quad_form_derivative(self, left_vecs, right_vecs):
         return self.base_lazy_tensor._quad_form_derivative(left_vecs, right_vecs)
 
-    def _solve(self, rhs, preconditioner, num_tridiag=None):
+    def _solve(self, rhs, preconditioner, num_tridiag=0):
         if num_tridiag:
             probe_vectors = rhs[..., :num_tridiag].detach()
             if torch.equal(probe_vectors, self.probe_vectors):
diff --git a/gpytorch/lazy/chol_lazy_tensor.py b/gpytorch/lazy/chol_lazy_tensor.py
index 990961dd..e8532554 100644
--- a/gpytorch/lazy/chol_lazy_tensor.py
+++ b/gpytorch/lazy/chol_lazy_tensor.py
@@ -33,14 +33,51 @@ class CholLazyTensor(RootLazyTensor):
             self._chol_diag_memo = self._chol.diagonal(dim1=-2, dim2=-1).clone()
         return self._chol_diag_memo
 
+    def _cholesky(self):
+        return self.root
+
+    def _solve(self, rhs, preconditioner, num_tridiag=0):
+        if num_tridiag:
+            return super()._solve(rhs, preconditioner, num_tridiag=num_tridiag)
+        else:
+            return self.root._cholesky_solve(rhs)
+
+    def inv_matmul(self, right_tensor, left_tensor=None):
+        with settings.fast_computations(solves=False):
+            return super().inv_matmul(right_tensor, left_tensor=left_tensor)
+
     def inv_quad_logdet(self, inv_quad_rhs=None, logdet=False, reduce_inv_quad=True):
+        if not self.is_square:
+            raise RuntimeError(
+                "inv_quad_logdet only operates on (batches of) square (positive semi-definite) LazyTensors. "
+                "Got a {} of size {}.".format(self.__class__.__name__, self.size())
+            )
+
+        if inv_quad_rhs is not None:
+            if self.dim() == 2 and inv_quad_rhs.dim() == 1:
+                if self.shape[-1] != inv_quad_rhs.numel():
+                    raise RuntimeError(
+                        "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                            self.shape, inv_quad_rhs.shape
+                        )
+                    )
+            elif self.dim() != inv_quad_rhs.dim():
+                raise RuntimeError(
+                    "LazyTensor (size={}) and right-hand-side Tensor (size={}) should have the same number "
+                    "of dimensions.".format(self.shape, inv_quad_rhs.shape)
+                )
+            elif self.batch_shape != inv_quad_rhs.shape[:-2] or self.shape[-1] != inv_quad_rhs.shape[-2]:
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                        self.shape, inv_quad_rhs.shape
+                    )
+                )
+
         inv_quad_term = None
         logdet_term = None
 
         if inv_quad_rhs is not None:
-            inv_quad_term, _ = super(CholLazyTensor, self).inv_quad_logdet(
-                inv_quad_rhs, logdet=False, reduce_inv_quad=reduce_inv_quad
-            )
+            inv_quad_term = self.inv_quad(inv_quad_rhs, reduce_inv_quad=reduce_inv_quad)
 
         if logdet:
             logdet_term = self._chol_diag.pow(2).log().sum(-1)
diff --git a/gpytorch/lazy/diag_lazy_tensor.py b/gpytorch/lazy/diag_lazy_tensor.py
index 695ad8a3..af4930db 100644
--- a/gpytorch/lazy/diag_lazy_tensor.py
+++ b/gpytorch/lazy/diag_lazy_tensor.py
@@ -3,7 +3,6 @@
 import torch
 from ..utils.memoize import cached
 from .lazy_tensor import LazyTensor
-from .root_lazy_tensor import RootLazyTensor
 
 
 class DiagLazyTensor(LazyTensor):
@@ -26,6 +25,13 @@ class DiagLazyTensor(LazyTensor):
 
         return AddedDiagLazyTensor(other, self)
 
+    @cached(name="cholesky")
+    def _cholesky(self):
+        return self.sqrt()
+
+    def _cholesky_solve(self, rhs):
+        return rhs / self._diag.pow(2)
+
     def _expand_batch(self, batch_shape):
         return self.__class__(self._diag.expand(*batch_shape, self._diag.size(-1)))
 
@@ -62,6 +68,12 @@ class DiagLazyTensor(LazyTensor):
             res = res.sum(-1)
         return (res,)
 
+    def _root_decomposition(self):
+        return self.sqrt()
+
+    def _root_inv_decomposition(self, initial_vectors=None):
+        return DiagLazyTensor(self._diag.reciprocal()).sqrt()
+
     def _size(self):
         return self._diag.shape + self._diag.shape[-1:]
 
@@ -134,12 +146,6 @@ class DiagLazyTensor(LazyTensor):
             return DiagLazyTensor(self._diag * other._diag)
         return super(DiagLazyTensor, self).matmul(other)
 
-    def root_decomposition(self):
-        return RootLazyTensor(self.sqrt())
-
-    def root_inv_decomposition(self):
-        return RootLazyTensor(DiagLazyTensor(self._diag.reciprocal()).sqrt())
-
     def sqrt(self):
         return DiagLazyTensor(self._diag.sqrt())
 
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 86a24f07..5ec7b84b 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -8,16 +8,17 @@ import gpytorch
 import torch
 
 from .. import settings
+from .. import utils
 from ..functions._inv_matmul import InvMatmul
+from ..functions._inv_quad import InvQuad
 from ..functions._inv_quad_log_det import InvQuadLogDet
 from ..functions._matmul import Matmul
 from ..functions._root_decomposition import RootDecomposition
-from ..utils import linear_cg
 from ..utils.broadcasting import _matmul_broadcast_shape, _mul_broadcast_shape
-from ..utils.cholesky import psd_safe_cholesky
+from ..utils.cholesky import psd_safe_cholesky, cholesky_solve
 from ..utils.deprecation import _deprecate_renamed_methods
 from ..utils.getitem import _noop_index, _convert_indices_to_tensors, _compute_getitem_size
-from ..utils.memoize import cached
+from ..utils.memoize import add_to_cache, cached
 from ..utils.qr import batch_qr
 from ..utils.svd import batch_svd
 from .lazy_tensor_representation_tree import LazyTensorRepresentationTree
@@ -379,6 +380,33 @@ class LazyTensor(ABC):
         """
         return self.diag()
 
+    @cached(name="cholesky")
+    def _cholesky(self):
+        """
+        (Optional) Cholesky-factorizes the LazyTensor
+
+        ..note::
+            This method is used as an internal helper. Calling this method directly is discouraged.
+
+        Returns:
+            (LazyTensor) Cholesky factor
+        """
+        from .non_lazy_tensor import NonLazyTensor
+        cholesky = psd_safe_cholesky(self.evaluate().double()).to(self.dtype)
+        return NonLazyTensor(cholesky)
+
+    def _cholesky_solve(self, rhs):
+        """
+        (Optional) Assuming that `self` is a Cholesky factor, computes the cholesky solve
+
+        ..note::
+            This method is used as an internal helper. Calling this method directly is discouraged.
+
+        Returns:
+            (LazyTensor) Cholesky factor
+        """
+        return cholesky_solve(rhs.double(), self.evaluate().double()).to(self.dtype)
+
     def _inv_matmul_preconditioner(self):
         """
         (Optional) define a preconditioner that can be used for linear systems, but not necessarily
@@ -452,7 +480,7 @@ class LazyTensor(ABC):
         if isinstance(self, NonLazyTensor) or isinstance(other, NonLazyTensor):
             return NonLazyTensor(self.evaluate() * other.evaluate())
         else:
-            left_lazy_tensor = self if self.root_decomposition_size() < other.root_decomposition_size() else other
+            left_lazy_tensor = self if self._root_decomposition_size() < other._root_decomposition_size() else other
             right_lazy_tensor = other if left_lazy_tensor is self else self
             return MulLazyTensor(left_lazy_tensor.root_decomposition(), right_lazy_tensor.root_decomposition())
 
@@ -521,8 +549,71 @@ class LazyTensor(ABC):
 
         return res
 
-    def _solve(self, rhs, preconditioner, num_tridiag=None):
-        return linear_cg(
+    def _root_decomposition(self):
+        """
+        Returns the (usually low-rank) root of a lazy tensor of a PSD matrix.
+
+        ..note::
+            This method is used internally by the related function
+            :func:`~gpytorch.lazy.LazyTensor.root_decomposition`, which does some additional work.
+            Calling this method directly is discouraged.
+
+        Returns:
+            (Tensor or LazyTensor): The root of the root decomposition
+        """
+        res, _ = RootDecomposition(
+            self.representation_tree(),
+            max_iter=self._root_decomposition_size(),
+            dtype=self.dtype,
+            device=self.device,
+            batch_shape=self.batch_shape,
+            matrix_shape=self.matrix_shape,
+        )(*self.representation())
+        return res
+
+    def _root_decomposition_size(self):
+        """
+        This is the inner size of the root decomposition.
+        This is primarily used to determine if it will be cheaper to compute a
+        different root or not
+        """
+        return settings.max_root_decomposition_size.value()
+
+    def _root_inv_decomposition(self, initial_vectors=None):
+        """
+        Returns the (usually low-rank) inverse root of a lazy tensor of a PSD matrix.
+
+        ..note::
+            This method is used internally by the related function
+            :func:`~gpytorch.lazy.LazyTensor.root_inv_decomposition`, which does some additional work.
+            Calling this method directly is discouraged.
+
+        Returns:
+            (Tensor or LazyTensor): The root of the inverse root decomposition
+        """
+        from .root_lazy_tensor import RootLazyTensor
+
+        roots, inv_roots = RootDecomposition(
+            self.representation_tree(),
+            max_iter=self._root_decomposition_size(),
+            dtype=self.dtype,
+            device=self.device,
+            batch_shape=self.batch_shape,
+            matrix_shape=self.matrix_shape,
+            root=True,
+            inverse=True,
+            initial_vectors=initial_vectors,
+        )(*self.representation())
+
+        if initial_vectors is not None and initial_vectors.size(-1) > 1:
+            add_to_cache(self, "root_decomposition", RootLazyTensor(roots[0]))
+        else:
+            add_to_cache(self, "root_decomposition", RootLazyTensor(roots))
+
+        return inv_roots
+
+    def _solve(self, rhs, preconditioner, num_tridiag=0):
+        return utils.linear_cg(
             self._matmul,
             rhs,
             n_tridiag=num_tridiag,
@@ -817,8 +908,37 @@ class LazyTensor(ABC):
         Returns:
             - tensor - tr( tensor^T (self)^{-1} tensor )
         """
-        res, _ = self.inv_quad_logdet(inv_quad_rhs=tensor, logdet=False, reduce_inv_quad=reduce_inv_quad)
-        return res
+        if not self.is_square:
+            raise RuntimeError(
+                "inv_quad 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 tensor.dim() == 1:
+            if self.shape[-1] != tensor.numel():
+                raise RuntimeError(
+                    "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                        self.shape, tensor.shape
+                    )
+                )
+        elif self.dim() != tensor.dim():
+            raise RuntimeError(
+                "LazyTensor (size={}) and right-hand-side Tensor (size={}) should have the same number "
+                "of dimensions.".format(self.shape, tensor.shape)
+            )
+        elif self.batch_shape != tensor.shape[:-2] or self.shape[-1] != tensor.shape[-2]:
+            raise RuntimeError(
+                "LazyTensor (size={}) cannot be multiplied with right-hand-side Tensor (size={}).".format(
+                    self.shape, tensor.shape
+                )
+            )
+
+        args = (tensor,) + self.representation()
+        inv_quad_term = InvQuad(representation_tree=self.representation_tree())(*args)
+
+        if reduce_inv_quad:
+            inv_quad_term = inv_quad_term.sum(-1)
+        return inv_quad_term
 
     def inv_quad_logdet(self, inv_quad_rhs=None, logdet=False, reduce_inv_quad=True):
         """
@@ -833,6 +953,15 @@ class LazyTensor(ABC):
             - scalar - tr( tensor^T (self)^{-1} tensor )
             - scalar - log determinant
         """
+        # Special case: use Cholesky to compute these terms
+        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())
+            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
+        # See NeurIPS 2018 paper: https://arxiv.org/abs/1809.11165
         if not self.is_square:
             raise RuntimeError(
                 "inv_quad_logdet only operates on (batches of) square (positive semi-definite) LazyTensors. "
@@ -1125,6 +1254,7 @@ class LazyTensor(ABC):
         This can be used for sampling from a Gaussian distribution, or for obtaining a
         low-rank version of a matrix
         """
+        from .chol_lazy_tensor import CholLazyTensor
         from .root_lazy_tensor import RootLazyTensor
 
         if not self.is_square:
@@ -1134,27 +1264,22 @@ class LazyTensor(ABC):
             )
 
         if (
-            self.matrix_shape.numel() <= settings.max_cholesky_numel.value()
+            self.size(-1) <= settings.max_cholesky_size.value()
             or settings.fast_computations.covar_root_decomposition.off()
         ):
             try:
-                return RootLazyTensor(psd_safe_cholesky(self.evaluate()))
+                res = self._cholesky()
+                return CholLazyTensor(res)
+
             except RuntimeError as e:
                 warnings.warn(
                     "Runtime Error when computing Cholesky decomposition: {}. Using RootDecomposition.".format(e)
                 )
 
-        res, _ = RootDecomposition(
-            self.representation_tree(),
-            max_iter=self.root_decomposition_size(),
-            dtype=self.dtype,
-            device=self.device,
-            batch_shape=self.batch_shape,
-            matrix_shape=self.matrix_shape,
-        )(*self.representation())
+        res = self._root_decomposition()
         return RootLazyTensor(res)
 
-    @cached
+    @cached(name="root_inv_decomposition")
     def root_inv_decomposition(self, initial_vectors=None, test_vectors=None):
         """
         Returns a (usually low-rank) root decomposotion lazy tensor of a PSD matrix.
@@ -1192,22 +1317,7 @@ class LazyTensor(ABC):
                     )
                 )
 
-        roots, inv_roots = RootDecomposition(
-            self.representation_tree(),
-            max_iter=self.root_decomposition_size(),
-            dtype=self.dtype,
-            device=self.device,
-            batch_shape=self.batch_shape,
-            matrix_shape=self.matrix_shape,
-            root=True,
-            inverse=True,
-            initial_vectors=initial_vectors,
-        )(*self.representation())
-
-        if initial_vectors is not None and initial_vectors.size(-1) > 1:
-            self._memoize_cache["root_decomposition"] = RootLazyTensor(roots[0])
-        else:
-            self._memoize_cache["root_decomposition"] = RootLazyTensor(roots)
+        inv_roots = self._root_inv_decomposition(initial_vectors)
 
         # Choose the best of the inv_roots, if there were more than one initial vectors
         if initial_vectors is not None and initial_vectors.size(-1) > 1:
@@ -1241,14 +1351,6 @@ class LazyTensor(ABC):
 
         return RootLazyTensor(inv_root)
 
-    def root_decomposition_size(self):
-        """
-        This is the inner size of the root decomposition.
-        This is primarily used to determine if it will be cheaper to compute a
-        different root or not
-        """
-        return settings.max_root_decomposition_size.value()
-
     def size(self, val=None):
         """
         Returns the size of the resulting Tensor that the lazy tensor represents
diff --git a/gpytorch/lazy/mul_lazy_tensor.py b/gpytorch/lazy/mul_lazy_tensor.py
index 73438e7f..a08131de 100644
--- a/gpytorch/lazy/mul_lazy_tensor.py
+++ b/gpytorch/lazy/mul_lazy_tensor.py
@@ -22,8 +22,10 @@ class MulLazyTensor(LazyTensor):
         Args:
             - lazy_tensors (A list of LazyTensor) - A list of LazyTensor to multiplicate with.
         """
-        left_lazy_tensor = left_lazy_tensor.root_decomposition()
-        right_lazy_tensor = right_lazy_tensor.root_decomposition()
+        if not isinstance(left_lazy_tensor, RootLazyTensor):
+            left_lazy_tensor = left_lazy_tensor.root_decomposition()
+        if not isinstance(right_lazy_tensor, RootLazyTensor):
+            right_lazy_tensor = right_lazy_tensor.root_decomposition()
         super(MulLazyTensor, self).__init__(left_lazy_tensor, right_lazy_tensor)
         self.left_lazy_tensor = left_lazy_tensor
         self.right_lazy_tensor = right_lazy_tensor
diff --git a/gpytorch/lazy/root_lazy_tensor.py b/gpytorch/lazy/root_lazy_tensor.py
index c523d576..c9e12b8c 100644
--- a/gpytorch/lazy/root_lazy_tensor.py
+++ b/gpytorch/lazy/root_lazy_tensor.py
@@ -69,6 +69,12 @@ class RootLazyTensor(LazyTensor):
             deriv.append(item_part_1 + item_part_2)
         return tuple(deriv)
 
+    def _root_decomposition(self):
+        return self.root
+
+    def _root_decomposition_size(self):
+        return self.root.size(-1)
+
     def _size(self):
         return torch.Size((*self.root.batch_shape, self.root.size(-2), self.root.size(-2)))
 
@@ -85,9 +91,3 @@ class RootLazyTensor(LazyTensor):
     def evaluate(self):
         eval_root = self.root.evaluate()
         return torch.matmul(eval_root, eval_root.transpose(-1, -2))
-
-    def root_decomposition_size(self):
-        return self.root.size(-1)
-
-    def root_decomposition(self):
-        return self
diff --git a/gpytorch/lazy/zero_lazy_tensor.py b/gpytorch/lazy/zero_lazy_tensor.py
index d5d7863e..4c29b47c 100644
--- a/gpytorch/lazy/zero_lazy_tensor.py
+++ b/gpytorch/lazy/zero_lazy_tensor.py
@@ -53,6 +53,15 @@ class ZeroLazyTensor(LazyTensor):
     def _quad_form_derivative(self, left_vecs, right_vecs):
         raise RuntimeError("Backwards through a ZeroLazyTensor is not possible")
 
+    def _root_decomposition(self):
+        raise RuntimeError("ZeroLazyTensors are not positive definite!")
+
+    def _root_inv_decomposition(self, initial_vectors=None):
+        raise RuntimeError("ZeroLazyTensors are not positive definite!")
+
+    def _root_decomposition_size(self):
+        raise RuntimeError("ZeroLazyTensors are not positive definite!")
+
     def _size(self):
         return torch.Size(self.sizes)
 
@@ -144,15 +153,6 @@ class ZeroLazyTensor(LazyTensor):
         shape = _mul_broadcast_shape(self.shape, other.shape)
         return self.__class__(*shape, dtype=self._dtype, device=self._device)
 
-    def root_decomposition(self):
-        raise RuntimeError("ZeroLazyTensors are not positive definite!")
-
-    def root_inv_decomposition(self, initial_vectors=None, test_vectors=None):
-        raise RuntimeError("ZeroLazyTensors are not positive definite!")
-
-    def root_decomposition_size(self):
-        raise RuntimeError("ZeroLazyTensors are not positive definite!")
-
     def transpose(self, dim1, dim2):
         sizes = self.sizes.copy()
         tmp = sizes[dim1]
diff --git a/gpytorch/settings.py b/gpytorch/settings.py
index 9b18f204..cf2d5428 100644
--- a/gpytorch/settings.py
+++ b/gpytorch/settings.py
@@ -101,6 +101,21 @@ class _fast_log_prob(_feature_flag):
     _state = True
 
 
+class _fast_solves(_feature_flag):
+    r"""
+    This feature flag controls how to compute solves with positive definite matrices.
+    If set to True, solves are computed using preconditioned conjugate gradients.
+    If set to False, `log_prob` is computed using the Cholesky decomposition.
+
+    .. warning ::
+
+        Setting this to False will compute a complete Cholesky decomposition of covariance matrices.
+        This may be infeasible for GPs with structure covariance matrices.
+    """
+
+    _state = True
+
+
 class check_training_data(_feature_flag):
     """
     Check whether the correct training data is supplied in Exact GP training mode
@@ -235,13 +250,22 @@ class fast_computations(object):
         * If set to False,
             `log_prob` is computed using the Cholesky decomposition.
 
+    * :attr:`fast_solves`
+        This feature flag controls how GPyTorch computes the solves of positive-definite matrices.
+
+        * If set to True,
+            Solves are computed with preconditioned conjugate gradients.
+
+        * If set to False,
+            Solves are computed using the Cholesky decomposition.
+
     .. warning ::
 
         Setting this to False will compute a complete Cholesky decomposition of covariance matrices.
         This may be infeasible for GPs with structure covariance matrices.
 
-    By default, approximations are used for all of these functions. Setting any of them to False will use
-    exact computations instead.
+    By default, approximations are used for all of these functions (except for solves).
+    Setting any of them to False will use exact computations instead.
 
     See also:
         * :class:`gpytorch.settings.max_root_decomposition_size`
@@ -254,17 +278,22 @@ class fast_computations(object):
     """
     covar_root_decomposition = _fast_covar_root_decomposition
     log_prob = _fast_log_prob
+    solves = _fast_solves
 
-    def __init__(self, covar_root_decomposition=True, log_prob=True):
+    def __init__(self, covar_root_decomposition=True, log_prob=True, solves=True):
         self.covar_root_decomposition = _fast_covar_root_decomposition(covar_root_decomposition)
         self.log_prob = _fast_log_prob(log_prob)
+        self.solves = _fast_solves(solves)
 
     def __enter__(self):
         self.covar_root_decomposition.__enter__()
         self.log_prob.__enter__()
+        self.solves.__enter__()
 
     def __exit__(self, *args):
+        self.covar_root_decomposition.__exit__()
         self.log_prob.__exit__()
+        self.solves.__exit__()
         return False
 
 
@@ -325,14 +354,14 @@ class _use_eval_tolerance(_feature_flag):
     _state = False
 
 
-class max_cholesky_numel(_value_context):
+class max_cholesky_size(_value_context):
     """
-    If the number of elements of a LazyTensor is less than `max_cholesky_numel`,
-    then the `root_decomposition` of LazyTensor will use Cholesky rather than Lanczos.
-    Default: 256
+    If the size of of a LazyTensor is less than `max_cholesky_size`,
+    then `root_decomposition` and `inv_matmul` of LazyTensor will use Cholesky rather than Lanczos/CG.
+    Default: 128
     """
 
-    _global_value = 256
+    _global_value = 128
 
 
 class max_root_decomposition_size(_value_context):
diff --git a/gpytorch/utils/memoize.py b/gpytorch/utils/memoize.py
index a30aaf50..878d5e93 100644
--- a/gpytorch/utils/memoize.py
+++ b/gpytorch/utils/memoize.py
@@ -3,6 +3,25 @@
 import functools
 
 
+def add_to_cache(obj, name, val):
+    """Add a result to the cache of an object."""
+    if not hasattr(obj, "_memoize_cache"):
+        obj._memoize_cache = dict()
+    obj._memoize_cache[name] = val
+    return obj
+
+
+def get_from_cache(obj, name):
+    """Get an item from the cache."""
+    if not is_in_cache(obj, name):
+        raise RuntimeError("Object does not have item {} stored in cache.".format(name))
+    return obj._memoize_cache[name]
+
+
+def is_in_cache(obj, name):
+    return hasattr(obj, "_memoize_cache") and name in obj._memoize_cache
+
+
 def cached(method=None, name=None):
     """A decorator allowing for specifying the name of a cache, allowing it to be modified elsewhere."""
     if method is None:
@@ -10,12 +29,10 @@ def cached(method=None, name=None):
 
     @functools.wraps(method)
     def g(self, *args, **kwargs):
-        if not hasattr(self, "_memoize_cache"):
-            self._memoize_cache = dict()
         cache_name = name if name is not None else method
-        if cache_name not in self._memoize_cache:
-            self._memoize_cache[cache_name] = method(self, *args, **kwargs)
-        return self._memoize_cache[cache_name]
+        if not is_in_cache(self, cache_name):
+            add_to_cache(self, cache_name, method(self, *args, **kwargs))
+        return get_from_cache(self, cache_name)
 
     return g
 
diff --git a/gpytorch/variational/variational_strategy.py b/gpytorch/variational/variational_strategy.py
index 2b268029..84cf303e 100644
--- a/gpytorch/variational/variational_strategy.py
+++ b/gpytorch/variational/variational_strategy.py
@@ -3,7 +3,7 @@
 import math
 import torch
 from .. import beta_features, settings
-from ..lazy import DiagLazyTensor, CachedCGLazyTensor, PsdSumLazyTensor, RootLazyTensor
+from ..lazy import DiagLazyTensor, CachedCGLazyTensor, CholLazyTensor, PsdSumLazyTensor, RootLazyTensor
 from ..module import Module
 from ..distributions import MultivariateNormal
 from ..utils.memoize import cached
@@ -130,6 +130,12 @@ class VariationalStrategy(Module):
             data_data_covar = full_covar[..., num_induc:, num_induc:]
             root_variational_covar = variational_dist.lazy_covariance_matrix.root_decomposition().root.evaluate()
 
+            # 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())
+                cholesky = True
+
             # Cache the CG results
             # For now: run variational inference without a preconditioner
             # The preconditioner screws things up for some reason
@@ -139,8 +145,8 @@ class VariationalStrategy(Module):
                 with torch.no_grad():
                     eager_rhs = torch.cat([left_tensors, induc_data_covar], -1)
                     solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
-                        induc_induc_covar, eager_rhs.detach(), logdet_terms=self.training,
-                        include_tmats=(not settings.skip_logdet_forward.on())
+                        induc_induc_covar, eager_rhs.detach(), logdet_terms=(not cholesky),
+                        include_tmats=(not settings.skip_logdet_forward.on() and not cholesky)
                     )
                     eager_rhss = [
                         eager_rhs.detach(), eager_rhs[..., left_tensors.size(-1):].detach(),
diff --git a/gpytorch/variational/whitened_variational_strategy.py b/gpytorch/variational/whitened_variational_strategy.py
index 3cf1d36a..8c137723 100644
--- a/gpytorch/variational/whitened_variational_strategy.py
+++ b/gpytorch/variational/whitened_variational_strategy.py
@@ -5,7 +5,8 @@ import torch
 from .. import settings, beta_features
 from .variational_strategy import VariationalStrategy
 from ..utils.memoize import cached
-from ..lazy import RootLazyTensor, MatmulLazyTensor, CachedCGLazyTensor, DiagLazyTensor, BatchRepeatLazyTensor
+from ..lazy import RootLazyTensor, MatmulLazyTensor, CholLazyTensor, \
+    CachedCGLazyTensor, DiagLazyTensor, BatchRepeatLazyTensor
 from ..distributions import MultivariateNormal
 
 
@@ -117,6 +118,12 @@ class WhitenedVariationalStrategy(VariationalStrategy):
             induc_data_covar = full_covar[..., :num_induc, num_induc:].evaluate()
             data_data_covar = full_covar[..., num_induc:, num_induc:]
 
+            # 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())
+                cholesky = True
+
             # Cache the CG results
             # Do not use preconditioning for whitened VI, as it does not seem to improve performance.
             with settings.max_preconditioner_size(0):
@@ -125,8 +132,8 @@ class WhitenedVariationalStrategy(VariationalStrategy):
                     solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
                         induc_induc_covar,
                         eager_rhs.detach(),
-                        logdet_terms=self.training,
-                        include_tmats=(not settings.skip_logdet_forward.on()),
+                        logdet_terms=(not cholesky),
+                        include_tmats=(not settings.skip_logdet_forward.on() and not cholesky),
                     )
                     eager_rhss = [eager_rhs.detach()]
                     solves = [solve.detach()]
diff --git a/readthedocs.yml b/readthedocs.yml
deleted file mode 100644
index 5f759dab..00000000
--- a/readthedocs.yml
+++ /dev/null
@@ -1,4 +0,0 @@
-build:
-  image: latest
-python:
-  version: 3.6
\ No newline at end of file


Test patch:
diff --git a/test/_base_test_case.py b/test/_base_test_case.py
new file mode 100644
index 00000000..c9cc5915
--- /dev/null
+++ b/test/_base_test_case.py
@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+
+from abc import ABC
+import os
+import random
+import torch
+
+
+class BaseTestCase(ABC):
+    def setUp(self):
+        if hasattr(self.__class__, "seed"):
+            seed = self.__class__.seed
+            if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
+                self.rng_state = torch.get_rng_state()
+                torch.manual_seed(seed)
+                if torch.cuda.is_available():
+                    torch.cuda.manual_seed_all(seed)
+                random.seed(seed)
+
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
+    def assertAllClose(self, tensor1, tensor2, rtol=1e-4, atol=1e-5, equal_nan=False):
+        if not tensor1.shape == tensor2.shape:
+            raise ValueError(f"tensor1 ({tensor1.shape}) and tensor2 ({tensor2.shape}) do not have the same shape.")
+
+        if torch.allclose(tensor1, tensor2, rtol=rtol, atol=atol, equal_nan=equal_nan):
+            return True
+
+        if not equal_nan:
+            if not torch.equal(tensor1, tensor1):
+                raise AssertionError(f"tensor1 ({tensor1.shape}) contains NaNs")
+            if not torch.equal(tensor2, tensor2):
+                raise AssertionError(f"tensor2 ({tensor2.shape}) contains NaNs")
+
+        rtol_diff = (torch.abs(tensor1 - tensor2) / torch.abs(tensor2)).view(-1)
+        rtol_diff = rtol_diff[torch.isfinite(rtol_diff)]
+        rtol_max = rtol_diff.max().item()
+
+        atol_diff = (torch.abs(tensor1 - tensor2) - torch.abs(tensor2).mul(rtol)).view(-1)
+        atol_diff = atol_diff[torch.isfinite(atol_diff)]
+        atol_max = atol_diff.max().item()
+
+        raise AssertionError(
+            f"tensor1 ({tensor1.shape}) and tensor2 ({tensor2.shape}) are not close enough. \n"
+            f"max rtol: {rtol_max:0.8f}\t\tmax atol: {atol_max:0.8f}"
+        )
diff --git a/test/examples/test_svgp_gp_regression.py b/test/examples/test_svgp_gp_regression.py
index fecf4e00..ffff7a3f 100644
--- a/test/examples/test_svgp_gp_regression.py
+++ b/test/examples/test_svgp_gp_regression.py
@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
 
-import os
-import random
+import math
+import warnings
 import unittest
-from math import pi
+from unittest.mock import MagicMock, patch
 from test._utils import least_used_cuda_device
 
 import gpytorch
@@ -12,11 +12,12 @@ from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.models import AbstractVariationalGP
 from gpytorch.variational import CholeskyVariationalDistribution, VariationalStrategy
 from torch import optim
+from .._base_test_case import BaseTestCase
 
 
 def train_data(cuda=False):
     train_x = torch.linspace(0, 1, 260)
-    train_y = torch.cos(train_x * (2 * pi))
+    train_y = torch.cos(train_x * (2 * math.pi))
     if cuda:
         return train_x.cuda(), train_y.cuda()
     else:
@@ -42,37 +43,43 @@ class SVGPRegressionModel(AbstractVariationalGP):
         return latent_pred
 
 
-class TestSVGPRegression(unittest.TestCase):
-    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 TestSVGPRegression(BaseTestCase, unittest.TestCase):
+    seed = 0
 
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
-
-    def test_regression_error(self, skip_logdet_forward=False):
-        train_x, train_y = train_data()
+    def test_regression_error(self, cuda=False, skip_logdet_forward=False, cholesky=False):
+        train_x, train_y = train_data(cuda=cuda)
         likelihood = GaussianLikelihood()
         model = SVGPRegressionModel(torch.linspace(0, 1, 25))
         mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=len(train_y))
+        if cuda:
+            likelihood = likelihood.cuda()
+            model = model.cuda()
+            mll = mll.cuda()
 
         # Find optimal model hyperparameters
         model.train()
         likelihood.train()
         optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
-        with gpytorch.settings.skip_logdet_forward(skip_logdet_forward):
-            for _ in range(170):
+
+        _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
+        with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), \
+                gpytorch.settings.skip_logdet_forward(skip_logdet_forward), \
+                warnings.catch_warnings(record=True) as w, \
+                patch("gpytorch.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
+            for _ in range(150):
                 optimizer.zero_grad()
                 output = model(train_x)
                 loss = -mll(output, train_y)
                 loss.backward()
                 optimizer.step()
 
+            # Make sure CG was called (or not), and no warnings were thrown
+            self.assertEqual(len(w), 0)
+            if cholesky:
+                self.assertFalse(linear_cg_mock.called)
+            else:
+                self.assertTrue(linear_cg_mock.called)
+
         for param in model.parameters():
             self.assertTrue(param.grad is not None)
             self.assertGreater(param.grad.norm().item(), 0)
@@ -90,40 +97,14 @@ class TestSVGPRegression(unittest.TestCase):
     def test_regression_error_skip_logdet_forward(self):
         return self.test_regression_error(skip_logdet_forward=True)
 
+    def test_regression_error_cholesky(self):
+        return self.test_regression_error(cholesky=True)
+
     def test_regression_error_cuda(self):
         if not torch.cuda.is_available():
             return
         with least_used_cuda_device():
-            train_x, train_y = train_data(cuda=True)
-            likelihood = GaussianLikelihood().cuda()
-            model = SVGPRegressionModel(torch.linspace(0, 1, 25)).cuda()
-            mll = gpytorch.mlls.VariationalMarginalLogLikelihood(likelihood, model, num_data=len(train_y))
-
-            # Find optimal model hyperparameters
-            model.train()
-            optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
-            optimizer.n_iter = 0
-            for _ in range(150):
-                optimizer.zero_grad()
-                output = model(train_x)
-                loss = -mll(output, train_y)
-                loss.backward()
-                optimizer.n_iter += 1
-                optimizer.step()
-
-            for param in model.parameters():
-                self.assertTrue(param.grad is not None)
-                self.assertGreater(param.grad.norm().item(), 0)
-            for param in likelihood.parameters():
-                self.assertTrue(param.grad is not None)
-                self.assertGreater(param.grad.norm().item(), 0)
-            optimizer.step()
-
-            # Set back to eval mode
-            model.eval()
-            test_preds = likelihood(model(train_x)).mean.squeeze()
-            mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
-            self.assertLess(mean_abs_error.item(), 1e-1)
+            return self.test_regression_error(cuda=True)
 
 
 if __name__ == "__main__":
diff --git a/test/examples/test_whitened_svgp_gp_regression.py b/test/examples/test_whitened_svgp_gp_regression.py
index a1c4bfee..9d86425c 100644
--- a/test/examples/test_whitened_svgp_gp_regression.py
+++ b/test/examples/test_whitened_svgp_gp_regression.py
@@ -1,10 +1,12 @@
 #!/usr/bin/env python3
 
-import os
-import random
+import math
 import unittest
-from math import pi
+import warnings
+import unittest
+from unittest.mock import MagicMock, patch
 from test._utils import least_used_cuda_device
+from .._base_test_case import BaseTestCase
 
 import gpytorch
 import torch
@@ -16,7 +18,7 @@ from torch import optim
 
 def train_data(cuda=False):
     train_x = torch.linspace(0, 1, 260)
-    train_y = torch.cos(train_x * (2 * pi))
+    train_y = torch.cos(train_x * (2 * math.pi))
     if cuda:
         return train_x.cuda(), train_y.cuda()
     else:
@@ -42,23 +44,14 @@ class SVGPRegressionModel(AbstractVariationalGP):
         return latent_pred
 
 
-class TestSVGPRegression(unittest.TestCase):
-    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)
-
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
+class TestSVGPRegression(BaseTestCase, unittest.TestCase):
+    seed = 0
 
-    def test_regression_error(self, skip_logdet_forward=False, cuda=False):
+    def test_regression_error(self, cuda=False, skip_logdet_forward=False, cholesky=False):
         train_x, train_y = train_data(cuda=cuda)
         likelihood = GaussianLikelihood()
-        model = SVGPRegressionModel(torch.linspace(0, 1, 25))
+        inducing_points = torch.linspace(0, 1, 25)
+        model = SVGPRegressionModel(inducing_points=inducing_points, learn_locs=False)
         if cuda:
             likelihood.cuda()
             model.cuda()
@@ -68,47 +61,12 @@ class TestSVGPRegression(unittest.TestCase):
         model.train()
         likelihood.train()
         optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
-        with gpytorch.settings.skip_logdet_forward(skip_logdet_forward):
-            for _ in range(170):
-                optimizer.zero_grad()
-                output = model(train_x)
-                loss = -mll(output, train_y)
-                loss.backward()
-                optimizer.step()
-
-        for param in model.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-        for param in likelihood.parameters():
-            self.assertTrue(param.grad is not None)
-            self.assertGreater(param.grad.norm().item(), 0)
-
-        # Set back to eval mode
-        model.eval()
-        likelihood.eval()
-        test_preds = likelihood(model(train_x)).mean.squeeze()
-        mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
-        self.assertLess(mean_abs_error.item(), 1e-1)
-
-    def test_regression_error_cuda(self):
-        if torch.cuda.is_available():
-            with least_used_cuda_device():
-                self.test_regression_error(skip_logdet_forward=False, cuda=True)
-
-    def test_regression_error_full(self, skip_logdet_forward=False, cuda=False):
-        train_x, train_y = train_data(cuda=cuda)
-        likelihood = GaussianLikelihood()
-        model = SVGPRegressionModel(inducing_points=train_x, learn_locs=False)
-        if cuda:
-            likelihood.cuda()
-            model.cuda()
-        mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=len(train_y))
 
-        # Find optimal model hyperparameters
-        model.train()
-        likelihood.train()
-        optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.01)
-        with gpytorch.settings.skip_logdet_forward(skip_logdet_forward):
+        _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
+        with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), \
+                gpytorch.settings.skip_logdet_forward(skip_logdet_forward), \
+                warnings.catch_warnings(record=True) as w, \
+                patch("gpytorch.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
             for _ in range(200):
                 optimizer.zero_grad()
                 output = model(train_x)
@@ -116,6 +74,13 @@ class TestSVGPRegression(unittest.TestCase):
                 loss.backward()
                 optimizer.step()
 
+            # Make sure CG was called (or not), and no warnings were thrown
+            self.assertEqual(len(w), 0)
+            if cholesky:
+                self.assertFalse(linear_cg_mock.called)
+            else:
+                self.assertTrue(linear_cg_mock.called)
+
         for param in model.parameters():
             self.assertTrue(param.grad is not None)
             self.assertGreater(param.grad.norm().item(), 0)
@@ -130,14 +95,12 @@ class TestSVGPRegression(unittest.TestCase):
         mean_abs_error = torch.mean(torch.abs(train_y - test_preds) / 2)
         self.assertLess(mean_abs_error.item(), 1e-1)
 
-    def test_regression_error_full_cuda(self):
-        if torch.cuda.is_available():
-            with least_used_cuda_device():
-                self.test_regression_error_full(skip_logdet_forward=False, cuda=True)
-
     def test_regression_error_skip_logdet_forward(self):
         self.test_regression_error(skip_logdet_forward=True)
 
+    def test_regression_error_cholesky(self):
+        return self.test_regression_error(cholesky=True)
+
     def test_regression_error_skip_logdet_forward_cuda(self):
         if torch.cuda.is_available():
             with least_used_cuda_device():
diff --git a/test/functions/test_inv_quad.py b/test/functions/test_inv_quad.py
new file mode 100644
index 00000000..54205276
--- /dev/null
+++ b/test/functions/test_inv_quad.py
@@ -0,0 +1,167 @@
+#!/usr/bin/env python3
+
+import os
+import random
+import torch
+import unittest
+import gpytorch
+from gpytorch.lazy import NonLazyTensor
+
+
+class TestInvQuadNonBatch(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)
+
+        mat = torch.randn(4, 4)
+        mat = mat @ mat.transpose(-1, -2)
+        mat.div_(5).add_(torch.eye(4))
+        vecs = torch.randn(5, 4, 6)
+        vec = torch.randn(4)
+        vecs = torch.randn(4, 8)
+        self.mat = mat.detach().clone().requires_grad_(True)
+        self.mat_clone = mat.detach().clone().requires_grad_(True)
+        self.vec = vec.detach().clone().requires_grad_(True)
+        self.vec_clone = vec.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+
+    def test_inv_quad_vector(self):
+        # Forward pass
+        actual_inv_quad = self.mat_clone.inverse().matmul(self.vec_clone).mul(self.vec_clone).sum()
+        with gpytorch.settings.num_trace_samples(1000):
+            non_lazy_tsr = NonLazyTensor(self.mat)
+            res_inv_quad = non_lazy_tsr.inv_quad(self.vec)
+
+        self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
+
+        # Backward
+        actual_inv_quad.backward()
+        res_inv_quad.backward(retain_graph=True)
+
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vec_clone.grad - self.vec.grad).abs()).item(), 1e-1)
+
+    def test_inv_quad_many_vectors(self):
+        # Forward pass
+        actual_inv_quad = self.mat_clone.inverse().matmul(self.vecs_clone).mul(self.vecs_clone).sum()
+        with gpytorch.settings.num_trace_samples(1000):
+            non_lazy_tsr = NonLazyTensor(self.mat)
+            res_inv_quad = non_lazy_tsr.inv_quad(self.vecs)
+        self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
+
+        # Backward
+        actual_inv_quad.backward()
+        res_inv_quad.backward(retain_graph=True)
+
+        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+
+
+class TestInvQuadBatch(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)
+
+        mats = torch.randn(5, 4, 4)
+        mats = mats @ mats.transpose(-1, -2)
+        mats.div_(5).add_(torch.eye(4).unsqueeze_(0))
+        vecs = torch.randn(5, 4, 6)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_clone = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+
+    def test_inv_quad_many_vectors(self):
+        # Forward pass
+        actual_inv_quad = (
+            torch.cat([mat.inverse().unsqueeze(0) for mat in self.mats_clone])
+            .matmul(self.vecs_clone)
+            .mul(self.vecs_clone)
+            .sum(2)
+            .sum(1)
+        )
+        with gpytorch.settings.num_trace_samples(2000):
+            non_lazy_tsr = NonLazyTensor(self.mats)
+            res_inv_quad = non_lazy_tsr.inv_quad(self.vecs)
+
+        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
+        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
+
+        # Backward
+        inv_quad_grad_output = torch.randn(5, dtype=torch.float)
+        actual_inv_quad.backward(gradient=inv_quad_grad_output)
+        res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
+
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+
+
+class TestInvQuadMultiBatch(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, 3, 4, 4)
+        mats = mats @ mats.transpose(-1, -2)
+        mats.div_(5).add_(torch.eye(4).view(1, 1, 4, 4))
+        vecs = torch.randn(2, 3, 4, 6)
+        self.mats = mats.detach().clone().requires_grad_(True)
+        self.mats_clone = mats.detach().clone().requires_grad_(True)
+        self.vecs = vecs.detach().clone().requires_grad_(True)
+        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+
+    def test_inv_quad_many_vectors(self):
+        # Forward pass
+        flattened_mats = self.mats_clone.view(-1, *self.mats_clone.shape[-2:])
+        actual_inv_quad = (
+            torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats])
+            .view(self.mats_clone.shape)
+            .matmul(self.vecs_clone)
+            .mul(self.vecs_clone)
+            .sum(-2)
+            .sum(-1)
+        )
+
+        with gpytorch.settings.num_trace_samples(2000):
+            non_lazy_tsr = NonLazyTensor(self.mats)
+            res_inv_quad = non_lazy_tsr.inv_quad(self.vecs)
+
+        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
+        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
+
+        # Backward
+        inv_quad_grad_output = torch.randn(2, 3, dtype=torch.float)
+        actual_inv_quad.backward(gradient=inv_quad_grad_output)
+        res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
+
+        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
+        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/functions/test_inv_quad_log_det.py b/test/functions/test_inv_quad_log_det.py
index e308adea..e312f1af 100644
--- a/test/functions/test_inv_quad_log_det.py
+++ b/test/functions/test_inv_quad_log_det.py
@@ -1,346 +1,109 @@
 #!/usr/bin/env python3
 
-import os
-import random
 import torch
 import unittest
-import gpytorch
-from gpytorch.lazy import NonLazyTensor
+from unittest.mock import MagicMock, patch
 
+import gpytorch
+from gpytorch.lazy import RootLazyTensor
+from .._base_test_case import BaseTestCase
+
+
+class TestInvQuadLogDetNonBatch(BaseTestCase, unittest.TestCase):
+    seed = 0
+    matrix_shape = torch.Size((4, 4))
+
+    def _test_inv_quad_logdet(self, inv_quad_rhs=None, logdet=False, improper_logdet=False):
+        # Set up
+        mat = torch.randn(*self.__class__.matrix_shape).requires_grad_(True)
+        mat_clone = mat.detach().clone().requires_grad_(True)
+
+        if inv_quad_rhs is not None:
+            inv_quad_rhs.requires_grad_(True)
+            inv_quad_rhs_clone = inv_quad_rhs.detach().clone().requires_grad_(True)
+
+        # Compute actual values
+        actual_tensor = mat_clone @ mat_clone.transpose(-1, -2)
+        if inv_quad_rhs is not None:
+            actual_inv_quad = actual_tensor.inverse().matmul(inv_quad_rhs_clone).mul(inv_quad_rhs_clone)
+            actual_inv_quad = actual_inv_quad.sum([-1, -2]) if inv_quad_rhs.dim() >= 2 else actual_inv_quad.sum()
+        if logdet:
+            flattened_tensor = actual_tensor.view(-1, *actual_tensor.shape[-2:])
+            logdets = torch.cat([mat.logdet().unsqueeze(0) for mat in flattened_tensor])
+            if actual_tensor.dim() > 2:
+                actual_logdet = logdets.view(*actual_tensor.shape[:-2])
+            else:
+                actual_logdet = logdets.squeeze()
+
+        # Compute values with LazyTensor
+        _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
+        with gpytorch.settings.num_trace_samples(2000), \
+                gpytorch.settings.max_cholesky_size(0), \
+                gpytorch.settings.cg_tolerance(1e-5), \
+                gpytorch.settings.skip_logdet_forward(improper_logdet), \
+                patch("gpytorch.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
+            lazy_tensor = RootLazyTensor(mat)
+            res_inv_quad, res_logdet = lazy_tensor.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet)
+
+        # Compare forward pass
+        if inv_quad_rhs is not None:
+            self.assertAllClose(res_inv_quad, actual_inv_quad, rtol=1e-2)
+        if logdet:
+            if improper_logdet:
+                self.assertAlmostEqual(res_logdet.norm().item(), 0)
+            else:
+                self.assertAllClose(res_logdet, actual_logdet, rtol=1e-1, atol=2e-1)
 
-class TestInvQuadLogDetNonBatch(unittest.TestCase):
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
+        # Backward
+        if inv_quad_rhs is not None:
+            actual_inv_quad.sum().backward(retain_graph=True)
+            res_inv_quad.sum().backward(retain_graph=True)
+        if logdet:
+            actual_logdet.sum().backward()
+            res_logdet.sum().backward()
 
-    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)
+        self.assertAllClose(mat_clone.grad, mat.grad, rtol=1e-1, atol=2e-1)
+        if inv_quad_rhs is not None:
+            self.assertAllClose(inv_quad_rhs.grad, inv_quad_rhs_clone.grad, rtol=1e-2)
 
-        mat = torch.randn(4, 4)
-        mat = mat @ mat.transpose(-1, -2)
-        mat.div_(5).add_(torch.eye(4))
-        vecs = torch.randn(5, 4, 6)
-        vec = torch.randn(4)
-        vecs = torch.randn(4, 8)
-        self.mat = mat.detach().clone().requires_grad_(True)
-        self.mat_clone = mat.detach().clone().requires_grad_(True)
-        self.vec = vec.detach().clone().requires_grad_(True)
-        self.vec_clone = vec.detach().clone().requires_grad_(True)
-        self.vecs = vecs.detach().clone().requires_grad_(True)
-        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+        # Make sure CG was called
+        self.assertTrue(linear_cg_mock.called)
 
     def test_inv_quad_logdet_vector(self):
-        # Forward pass
-        actual_inv_quad = self.mat_clone.inverse().matmul(self.vec_clone).mul(self.vec_clone).sum()
-        actual_logdet = self.mat_clone.logdet()
-        with gpytorch.settings.num_trace_samples(1000):
-            non_lazy_tsr = NonLazyTensor(self.mat)
-            res_inv_quad, res_logdet = non_lazy_tsr.inv_quad_logdet(inv_quad_rhs=self.vec, logdet=True)
-
-        self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
-        self.assertAlmostEqual(res_logdet.item(), actual_logdet.item(), places=1)
-
-        # Backward
-        actual_inv_quad.backward()
-        actual_logdet.backward()
-        res_inv_quad.backward(retain_graph=True)
-        res_logdet.backward()
-
-        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((self.vec_clone.grad - self.vec.grad).abs()).item(), 1e-1)
+        rhs = torch.randn(self.matrix_shape[-1])
+        self._test_inv_quad_logdet(inv_quad_rhs=rhs, logdet=True)
 
     def test_inv_quad_only_vector(self):
-        # Forward pass
-        res = NonLazyTensor(self.mat).inv_quad(self.vec)
-        actual = self.mat_clone.inverse().matmul(self.vec_clone).mul(self.vec_clone).sum()
-        self.assertAlmostEqual(res.item(), actual.item(), places=1)
-
-        # Backward
-        actual.backward()
-        res.backward()
-
-        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-3)
-        self.assertLess(torch.max((self.vec_clone.grad - self.vec.grad).abs()).item(), 1e-3)
+        rhs = torch.randn(self.matrix_shape[-1])
+        self._test_inv_quad_logdet(inv_quad_rhs=rhs, logdet=False)
 
     def test_inv_quad_logdet_many_vectors(self):
-        # Forward pass
-        actual_inv_quad = self.mat_clone.inverse().matmul(self.vecs_clone).mul(self.vecs_clone).sum()
-        actual_logdet = self.mat_clone.logdet()
-        with gpytorch.settings.num_trace_samples(1000):
-            non_lazy_tsr = NonLazyTensor(self.mat)
-            res_inv_quad, res_logdet = non_lazy_tsr.inv_quad_logdet(inv_quad_rhs=self.vecs, logdet=True)
-        self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
-        self.assertAlmostEqual(res_logdet.item(), actual_logdet.item(), places=1)
-
-        # Backward
-        actual_inv_quad.backward()
-        actual_logdet.backward()
-        res_inv_quad.backward(retain_graph=True)
-        res_logdet.backward()
-
-        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+        rhs = torch.randn(*self.matrix_shape[:-1], 5)
+        self._test_inv_quad_logdet(inv_quad_rhs=rhs, logdet=True)
 
     def test_inv_quad_logdet_many_vectors_improper(self):
-        # Forward pass
-        actual_inv_quad = self.mat_clone.inverse().matmul(self.vecs_clone).mul(self.vecs_clone).sum()
-        actual_logdet = self.mat_clone.logdet()
-        with gpytorch.settings.num_trace_samples(1000), gpytorch.settings.skip_logdet_forward(True):
-            non_lazy_tsr = NonLazyTensor(self.mat)
-            res_inv_quad, res_logdet = non_lazy_tsr.inv_quad_logdet(inv_quad_rhs=self.vecs, logdet=True)
-        self.assertAlmostEqual(res_inv_quad.item(), actual_inv_quad.item(), places=1)
-        self.assertAlmostEqual(res_logdet.item(), torch.zeros(non_lazy_tsr.batch_shape), places=1)
-
-        # Backward
-        actual_inv_quad.backward()
-        actual_logdet.backward()
-        res_inv_quad.backward(retain_graph=True)
-        res_logdet.backward()
-
-        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
+        rhs = torch.randn(*self.matrix_shape[:-1], 5)
+        self._test_inv_quad_logdet(inv_quad_rhs=rhs, logdet=True, improper_logdet=True)
 
     def test_inv_quad_only_many_vectors(self):
-        # Forward pass
-        res = NonLazyTensor(self.mat).inv_quad(self.vecs)
-        actual = self.mat_clone.inverse().matmul(self.vecs_clone).mul(self.vecs_clone).sum()
-        self.assertAlmostEqual(res.item(), actual.item(), places=1)
+        rhs = torch.randn(*self.matrix_shape[:-1], 5)
+        self._test_inv_quad_logdet(inv_quad_rhs=rhs, logdet=False)
 
-        # Backward
-        actual.backward()
-        res.backward()
-
-        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-3)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-3)
-
-    def test_logdet_only(self):
-        # Forward pass
-        with gpytorch.settings.num_trace_samples(1000):
-            res = NonLazyTensor(self.mat).logdet()
-        actual = self.mat_clone.logdet()
-        self.assertAlmostEqual(res.item(), actual.item(), places=1)
-
-        # Backward
-        actual.backward()
-        res.backward()
-        self.assertLess(torch.max((self.mat_clone.grad - self.mat.grad).abs()).item(), 1e-1)
-
-
-class TestInvQuadLogDetBatch(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)
 
-        mats = torch.randn(5, 4, 4)
-        mats = mats @ mats.transpose(-1, -2)
-        mats.div_(5).add_(torch.eye(4).unsqueeze_(0))
-        vecs = torch.randn(5, 4, 6)
-        self.mats = mats.detach().clone().requires_grad_(True)
-        self.mats_clone = mats.detach().clone().requires_grad_(True)
-        self.vecs = vecs.detach().clone().requires_grad_(True)
-        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
+class TestInvQuadLogDetBatch(TestInvQuadLogDetNonBatch):
+    seed = 0
+    matrix_shape = torch.Size((3, 4, 4))
 
-    def test_inv_quad_logdet_many_vectors(self):
-        # Forward pass
-        actual_inv_quad = (
-            torch.cat([mat.inverse().unsqueeze(0) for mat in self.mats_clone])
-            .matmul(self.vecs_clone)
-            .mul(self.vecs_clone)
-            .sum(2)
-            .sum(1)
-        )
-        actual_logdet = torch.cat([mat.logdet().unsqueeze(0) for mat in self.mats_clone])
-        with gpytorch.settings.num_trace_samples(2000):
-            non_lazy_tsr = NonLazyTensor(self.mats)
-            res_inv_quad, res_logdet = non_lazy_tsr.inv_quad_logdet(inv_quad_rhs=self.vecs, logdet=True)
-
-        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
-        self.assertEqual(res_logdet.shape, actual_logdet.shape)
-        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((res_logdet - actual_logdet).abs()).item(), 1e-1)
-
-        # Backward
-        inv_quad_grad_output = torch.randn(5, dtype=torch.float)
-        logdet_grad_output = torch.randn(5, dtype=torch.float)
-        actual_inv_quad.backward(gradient=inv_quad_grad_output)
-        actual_logdet.backward(gradient=logdet_grad_output)
-        res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
-        res_logdet.backward(gradient=logdet_grad_output)
-
-        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
-
-    def test_inv_quad_logdet_many_vectors_improper(self):
-        # Forward pass
-        actual_inv_quad = (
-            torch.cat([mat.inverse().unsqueeze(0) for mat in self.mats_clone])
-            .matmul(self.vecs_clone)
-            .mul(self.vecs_clone)
-            .sum(2)
-            .sum(1)
-        )
-        actual_logdet = torch.cat([mat.logdet().unsqueeze(0) for mat in self.mats_clone])
-        with gpytorch.settings.num_trace_samples(2000), gpytorch.settings.skip_logdet_forward(True):
-            non_lazy_tsr = NonLazyTensor(self.mats)
-            res_inv_quad, res_logdet = non_lazy_tsr.inv_quad_logdet(inv_quad_rhs=self.vecs, logdet=True)
-
-        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
-        self.assertEqual(res_logdet.shape, actual_logdet.shape)
-        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
-        self.assertLess(torch.max(res_logdet.abs()).item(), 1e-1)
-
-        # Backward
-        inv_quad_grad_output = torch.randn(5, dtype=torch.float)
-        logdet_grad_output = torch.randn(5, dtype=torch.float)
-        actual_inv_quad.backward(gradient=inv_quad_grad_output)
-        actual_logdet.backward(gradient=logdet_grad_output)
-        res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
-        res_logdet.backward(gradient=logdet_grad_output)
-
-        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
-
-    def test_inv_quad_only_many_vectors(self):
-        # Forward pass
-        res = NonLazyTensor(self.mats).inv_quad(self.vecs)
-        actual = (
-            torch.cat([mat.inverse().unsqueeze(0) for mat in self.mats_clone])
-            .matmul(self.vecs_clone)
-            .mul(self.vecs_clone)
-            .sum(2)
-            .sum(1)
-        )
-        self.assertEqual(res.shape, actual.shape)
-        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
-
-        # Backward
-        actual.sum().backward()
-        res.sum().backward()
-
-        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-3)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-3)
-
-    def test_logdet_only(self):
-        # Forward pass
-        with gpytorch.settings.num_trace_samples(2000):
-            res = NonLazyTensor(self.mats).logdet()
-        actual = torch.cat([mat.logdet().unsqueeze(0) for mat in self.mats_clone])
-        self.assertEqual(res.shape, actual.shape)
-        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
-
-        # Backward
-        grad_output = torch.randn(5)
-        actual.backward(gradient=grad_output)
-        res.backward(gradient=grad_output)
-        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
-
-
-class TestInvQuadLogDetMultiBatch(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, 3, 4, 4)
-        mats = mats @ mats.transpose(-1, -2)
-        mats.div_(5).add_(torch.eye(4).view(1, 1, 4, 4))
-        vecs = torch.randn(2, 3, 4, 6)
-        self.mats = mats.detach().clone().requires_grad_(True)
-        self.mats_clone = mats.detach().clone().requires_grad_(True)
-        self.vecs = vecs.detach().clone().requires_grad_(True)
-        self.vecs_clone = vecs.detach().clone().requires_grad_(True)
-
-    def test_inv_quad_logdet_many_vectors(self):
-        # Forward pass
-        flattened_mats = self.mats_clone.view(-1, *self.mats_clone.shape[-2:])
-        actual_inv_quad = (
-            torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats])
-            .view(self.mats_clone.shape)
-            .matmul(self.vecs_clone)
-            .mul(self.vecs_clone)
-            .sum(-2)
-            .sum(-1)
-        )
-        actual_logdet = torch.cat([mat.logdet().unsqueeze(0) for mat in flattened_mats])
-        actual_logdet = actual_logdet.view(self.mats_clone.shape[:-2])
-
-        with gpytorch.settings.num_trace_samples(2000):
-            non_lazy_tsr = NonLazyTensor(self.mats)
-            res_inv_quad, res_logdet = non_lazy_tsr.inv_quad_logdet(inv_quad_rhs=self.vecs, logdet=True)
-
-        self.assertEqual(res_inv_quad.shape, actual_inv_quad.shape)
-        self.assertEqual(res_logdet.shape, actual_logdet.shape)
-        self.assertLess(torch.max((res_inv_quad - actual_inv_quad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((res_logdet - actual_logdet).abs()).item(), 1e-1)
-
-        # Backward
-        inv_quad_grad_output = torch.randn(2, 3, dtype=torch.float)
-        logdet_grad_output = torch.randn(2, 3, dtype=torch.float)
-        actual_inv_quad.backward(gradient=inv_quad_grad_output)
-        actual_logdet.backward(gradient=logdet_grad_output)
-        res_inv_quad.backward(gradient=inv_quad_grad_output, retain_graph=True)
-        res_logdet.backward(gradient=logdet_grad_output)
-
-        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-1)
-
-    def test_inv_quad_only_many_vectors(self):
-        # Forward pass
-        res = NonLazyTensor(self.mats).inv_quad(self.vecs)
-        flattened_mats = self.mats_clone.view(-1, *self.mats_clone.shape[-2:])
-        actual = (
-            torch.cat([mat.inverse().unsqueeze(0) for mat in flattened_mats])
-            .view(self.mats_clone.shape)
-            .matmul(self.vecs_clone)
-            .mul(self.vecs_clone)
-            .sum(-2)
-            .sum(-1)
-        )
-        self.assertEqual(res.shape, actual.shape)
-        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
-
-        # Backward
-        actual.sum().backward()
-        res.sum().backward()
+    def test_inv_quad_logdet_vector(self):
+        pass
 
-        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-3)
-        self.assertLess(torch.max((self.vecs_clone.grad - self.vecs.grad).abs()).item(), 1e-3)
+    def test_inv_quad_only_vector(self):
+        pass
 
-    def test_logdet_only(self):
-        # Forward pass
-        with gpytorch.settings.num_trace_samples(2000):
-            res = NonLazyTensor(self.mats).logdet()
-        flattened_mats = self.mats_clone.view(-1, *self.mats_clone.shape[-2:])
-        actual = torch.cat([mat.logdet().unsqueeze(0) for mat in flattened_mats])
-        actual = actual.view(self.mats_clone.shape[:-2])
-        self.assertEqual(res.shape, actual.shape)
-        self.assertLess(torch.max((res - actual).abs()).item(), 1e-1)
 
-        # Backward
-        grad_output = torch.randn(2, 3)
-        actual.backward(gradient=grad_output)
-        res.backward(gradient=grad_output)
-        self.assertLess(torch.max((self.mats_clone.grad - self.mats.grad).abs()).item(), 1e-1)
+class TestInvQuadLogDetMultiBatch(TestInvQuadLogDetBatch):
+    seed = 0
+    matrix_shape = torch.Size((2, 3, 4, 4))
 
 
 if __name__ == "__main__":
diff --git a/test/lazy/_lazy_tensor_test_case.py b/test/lazy/_lazy_tensor_test_case.py
index 56579210..59bba24a 100644
--- a/test/lazy/_lazy_tensor_test_case.py
+++ b/test/lazy/_lazy_tensor_test_case.py
@@ -1,17 +1,16 @@
 #!/usr/bin/env python3
 
 import math
-import os
-import random
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from itertools import product, combinations
-from test._utils import approx_equal
+from unittest.mock import MagicMock, patch
 
 import gpytorch
 import torch
+from .._base_test_case import BaseTestCase
 
 
-class RectangularLazyTensorTestCase(ABC):
+class RectangularLazyTensorTestCase(BaseTestCase):
     @abstractmethod
     def create_lazy_tensor(self):
         raise NotImplementedError()
@@ -20,129 +19,66 @@ class RectangularLazyTensorTestCase(ABC):
     def evaluate_lazy_tensor(self):
         raise NotImplementedError()
 
-    def setUp(self):
-        if hasattr(self.__class__, "seed"):
-            seed = self.__class__.seed
-            if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
-                self.rng_state = torch.get_rng_state()
-                torch.manual_seed(seed)
-                if torch.cuda.is_available():
-                    torch.cuda.manual_seed_all(seed)
-                random.seed(seed)
-
-    def tearDown(self):
-        if hasattr(self, "rng_state"):
-            torch.set_rng_state(self.rng_state)
-
-    def test_matmul_vec(self):
+    def _test_matmul(self, rhs):
         lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
-
-        # We skip this test if we're dealing with batch LazyTensors
-        # They shouldn't multiply by a vec
-        if lazy_tensor.ndimension() > 2:
-            return
-
         lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
         evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
 
-        test_vector = torch.randn(lazy_tensor.size(-1))
-        res = lazy_tensor.matmul(test_vector)
-        actual = evaluated.matmul(test_vector)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+        res = lazy_tensor.matmul(rhs)
+        actual = evaluated.matmul(rhs)
+        self.assertAllClose(res, actual)
 
         grad = torch.randn_like(res)
         res.backward(gradient=grad)
         actual.backward(gradient=grad)
         for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
             if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
+                self.assertAllClose(arg.grad, arg_copy.grad, rtol=1e-3)
 
-    def test_matmul_matrix(self):
-        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)
+    def test_matmul_vec(self):
+        lazy_tensor = self.create_lazy_tensor()
+        rhs = torch.randn(lazy_tensor.size(-1))
 
-        test_vector = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
-        res = lazy_tensor.matmul(test_vector)
-        actual = evaluated.matmul(test_vector)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+        # We skip this test if we're dealing with batch LazyTensors
+        # They shouldn't multiply by a vec
+        if lazy_tensor.ndimension() > 2:
+            return
+        else:
+            return self._test_matmul(rhs)
 
-        grad = torch.randn_like(res)
-        res.backward(gradient=grad)
-        actual.backward(gradient=grad)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
+    def test_matmul_matrix(self):
+        lazy_tensor = self.create_lazy_tensor()
+        rhs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 4)
+        return self._test_matmul(rhs)
 
     def test_matmul_matrix_broadcast(self):
-        # Right hand size has one more batch dimension
-        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)
-        test_vector = torch.randn(3, *lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
-        res = lazy_tensor.matmul(test_vector)
-        actual = evaluated.matmul(test_vector)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+        lazy_tensor = self.create_lazy_tensor()
 
-        grad = torch.randn_like(res)
-        res.backward(gradient=grad)
-        actual.backward(gradient=grad)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
+        # Right hand size has one more batch dimension
+        batch_shape = torch.Size((3, *lazy_tensor.batch_shape))
+        rhs = torch.randn(*batch_shape, lazy_tensor.size(-1), 4)
+        self._test_matmul(rhs)
 
         if lazy_tensor.ndimension() > 2:
             # Right hand size has one fewer batch dimension
-            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)
-            test_vector = torch.randn(*lazy_tensor.batch_shape[1:], lazy_tensor.size(-1), 5)
-            res = lazy_tensor.matmul(test_vector)
-            actual = evaluated.matmul(test_vector)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-            grad = torch.randn_like(res)
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-                if arg_copy.grad is not None:
-                    self.assertLess(
-                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                    )
+            batch_shape = torch.Size(lazy_tensor.batch_shape[1:])
+            rhs = torch.randn(*batch_shape, lazy_tensor.size(-1), 4)
+            self._test_matmul(rhs)
 
             # Right hand size has a singleton dimension
-            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)
-            test_vector = torch.randn(*lazy_tensor.batch_shape[:-1], 1, lazy_tensor.size(-1), 5)
-            res = lazy_tensor.matmul(test_vector)
-            actual = evaluated.matmul(test_vector)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-            grad = torch.randn_like(res)
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-                if arg_copy.grad is not None:
-                    self.assertLess(
-                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                    )
+            batch_shape = torch.Size((*lazy_tensor.batch_shape[:-1], 1))
+            rhs = torch.randn(*batch_shape, lazy_tensor.size(-1), 4)
+            self._test_matmul(rhs)
 
     def test_constant_mul(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
-        self.assertTrue(approx_equal((lazy_tensor * 5).evaluate(), evaluated * 5))
+        self.assertAllClose((lazy_tensor * 5).evaluate(), evaluated * 5)
 
     def test_evaluate(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
-        self.assertTrue(approx_equal(lazy_tensor.evaluate(), evaluated))
+        self.assertAllClose(lazy_tensor.evaluate(), evaluated)
 
     def test_getitem(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -152,75 +88,59 @@ class RectangularLazyTensorTestCase(ABC):
         if lazy_tensor.ndimension() == 2:
             res = lazy_tensor[1]
             actual = evaluated[1]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[0:2].evaluate()
             actual = evaluated[0:2]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[:, 0:2].evaluate()
             actual = evaluated[:, 0:2]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[0:2, :].evaluate()
             actual = evaluated[0:2, :]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[..., 0:2].evaluate()
             actual = evaluated[..., 0:2]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[0:2, ...].evaluate()
             actual = evaluated[0:2, ...]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[..., 0:2, 2]
             actual = evaluated[..., 0:2, 2]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[0:2, ..., 2]
             actual = evaluated[0:2, ..., 2]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
 
         # Batch case
         else:
             res = lazy_tensor[1].evaluate()
             actual = evaluated[1]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[0:2].evaluate()
             actual = evaluated[0:2]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor[:, 0:2].evaluate()
             actual = evaluated[:, 0:2]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
 
             for batch_index in product([1, slice(0, 2, None)], repeat=(lazy_tensor.dim() - 2)):
                 res = lazy_tensor.__getitem__((*batch_index, slice(0, 1, None), slice(0, 2, None))).evaluate()
                 actual = evaluated.__getitem__((*batch_index, slice(0, 1, None), slice(0, 2, None)))
-                self.assertEqual(res.shape, actual.shape)
-                self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+                self.assertAllClose(res, actual)
                 res = lazy_tensor.__getitem__((*batch_index, 1, slice(0, 2, None)))
                 actual = evaluated.__getitem__((*batch_index, 1, slice(0, 2, None)))
-                self.assertEqual(res.shape, actual.shape)
-                self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+                self.assertAllClose(res, actual)
                 res = lazy_tensor.__getitem__((*batch_index, slice(1, None, None), 2))
                 actual = evaluated.__getitem__((*batch_index, slice(1, None, None), 2))
-                self.assertEqual(res.shape, actual.shape)
-                self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+                self.assertAllClose(res, actual)
 
             # Ellipsis
             res = lazy_tensor.__getitem__((Ellipsis, slice(1, None, None), 2))
             actual = evaluated.__getitem__((Ellipsis, slice(1, None, None), 2))
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = lazy_tensor.__getitem__((slice(1, None, None), Ellipsis, 2))
             actual = evaluated.__getitem__((slice(1, None, None), Ellipsis, 2))
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
 
     def test_getitem_tensor_index(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -230,28 +150,22 @@ class RectangularLazyTensorTestCase(ABC):
         if lazy_tensor.ndimension() == 2:
             index = (torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
             res, actual = lazy_tensor[index], evaluated[index]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             index = (torch.tensor([0, 0, 1, 2]), slice(None, None, None))
             res, actual = gpytorch.delazify(lazy_tensor[index]), evaluated[index]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]))
             res, actual = gpytorch.delazify(lazy_tensor[index]), evaluated[index]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             index = (torch.tensor([0, 0, 1, 2]), Ellipsis)
             res, actual = gpytorch.delazify(lazy_tensor[index]), evaluated[index]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             index = (Ellipsis, torch.tensor([0, 0, 1, 2]))
             res, actual = gpytorch.delazify(lazy_tensor[index]), evaluated[index]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             index = (Ellipsis, torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
             res, actual = lazy_tensor[index], evaluated[index]
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
 
         # Batch case
         else:
@@ -260,32 +174,26 @@ class RectangularLazyTensorTestCase(ABC):
             ):
                 index = (*batch_index, torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1]))
                 res, actual = lazy_tensor[index], evaluated[index]
-                self.assertEqual(res.shape, actual.shape)
-                self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+                self.assertAllClose(res, actual)
                 index = (*batch_index, torch.tensor([0, 1, 0, 2]), slice(None, None, None))
                 res, actual = gpytorch.delazify(lazy_tensor[index]), evaluated[index]
-                self.assertEqual(res.shape, actual.shape)
-                self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+                self.assertAllClose(res, actual)
                 index = (*batch_index, slice(None, None, None), torch.tensor([0, 1, 2, 1]))
                 res, actual = gpytorch.delazify(lazy_tensor[index]), evaluated[index]
-                self.assertEqual(res.shape, actual.shape)
-                self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+                self.assertAllClose(res, actual)
                 index = (*batch_index, slice(None, None, None), slice(None, None, None))
                 res, actual = lazy_tensor[index].evaluate(), evaluated[index]
-                self.assertEqual(res.shape, actual.shape)
-                self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+                self.assertAllClose(res, actual)
 
             # Ellipsis
             res = lazy_tensor.__getitem__((Ellipsis, torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1])))
             actual = evaluated.__getitem__((Ellipsis, torch.tensor([0, 1, 0, 2]), torch.tensor([1, 2, 0, 1])))
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
             res = gpytorch.delazify(
                 lazy_tensor.__getitem__((torch.tensor([0, 1, 0, 1]), Ellipsis, torch.tensor([1, 2, 0, 1])))
             )
             actual = evaluated.__getitem__((torch.tensor([0, 1, 0, 1]), Ellipsis, torch.tensor([1, 2, 0, 1])))
-            self.assertEqual(res.shape, actual.shape)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 1e-1)
+            self.assertAllClose(res, actual)
 
     def test_permute(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -294,7 +202,7 @@ class RectangularLazyTensorTestCase(ABC):
             dims = torch.randperm(lazy_tensor.dim() - 2).tolist()
             res = lazy_tensor.permute(*dims, -2, -1).evaluate()
             actual = evaluated.permute(*dims, -2, -1)
-            self.assertTrue(approx_equal(res, actual))
+            self.assertAllClose(res, actual)
 
     def test_quad_form_derivative(self):
         lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
@@ -306,18 +214,18 @@ class RectangularLazyTensorTestCase(ABC):
         deriv_auto = gpytorch.lazy.LazyTensor._quad_form_derivative(lazy_tensor_clone, left_vecs, right_vecs)
 
         for dc, da in zip(deriv_custom, deriv_auto):
-            self.assertLess(torch.norm(dc - da), 1e-1)
+            self.assertAllClose(dc, da)
 
     def test_sum(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
 
-        self.assertTrue(approx_equal(lazy_tensor.sum(-1), evaluated.sum(-1)))
-        self.assertTrue(approx_equal(lazy_tensor.sum(-2), evaluated.sum(-2)))
+        self.assertAllClose(lazy_tensor.sum(-1), evaluated.sum(-1))
+        self.assertAllClose(lazy_tensor.sum(-2), evaluated.sum(-2))
         if lazy_tensor.ndimension() > 2:
-            self.assertTrue(approx_equal(lazy_tensor.sum(-3).evaluate(), evaluated.sum(-3)))
+            self.assertAllClose(lazy_tensor.sum(-3).evaluate(), evaluated.sum(-3))
         if lazy_tensor.ndimension() > 3:
-            self.assertTrue(approx_equal(lazy_tensor.sum(-4).evaluate(), evaluated.sum(-4)))
+            self.assertAllClose(lazy_tensor.sum(-4).evaluate(), evaluated.sum(-4))
 
     def test_transpose_batch(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -327,12 +235,91 @@ class RectangularLazyTensorTestCase(ABC):
             for i, j in combinations(range(lazy_tensor.dim() - 2), 2):
                 res = lazy_tensor.transpose(i, j).evaluate()
                 actual = evaluated.transpose(i, j)
-                self.assertTrue(torch.allclose(res, actual, rtol=1e-4, atol=1e-5))
+                self.assertAllClose(res, actual, rtol=1e-4, atol=1e-5)
 
 
 class LazyTensorTestCase(RectangularLazyTensorTestCase):
     should_test_sample = False
     skip_slq_tests = False
+    should_call_cg = True
+    should_call_lanczos = True
+
+    def _test_inv_matmul(self, rhs, lhs=None, cholesky=False):
+        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)
+
+        # Create a test right hand side and left hand side
+        rhs.requires_grad_(True)
+        rhs_copy = rhs.clone().detach().requires_grad_(True)
+        if lhs is not None:
+            lhs.requires_grad_(True)
+            lhs_copy = lhs.clone().detach().requires_grad_(True)
+
+        _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
+        with patch("gpytorch.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
+            with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), \
+                    gpytorch.settings.cg_tolerance(1e-4):
+                # Perform the inv_matmul
+                if lhs is not None:
+                    res = lazy_tensor.inv_matmul(rhs, lhs)
+                    actual = lhs_copy @ evaluated.inverse() @ rhs_copy
+                else:
+                    res = lazy_tensor.inv_matmul(rhs)
+                    actual = evaluated.inverse().matmul(rhs_copy)
+                self.assertAllClose(res, actual, rtol=0.02, atol=1e-5)
+
+                # Perform backward pass
+                grad = torch.randn_like(res)
+                res.backward(gradient=grad)
+                actual.backward(gradient=grad)
+                for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+                    if arg_copy.grad is not None:
+                        self.assertAllClose(arg.grad, arg_copy.grad, rtol=0.03, atol=1e-5)
+                self.assertAllClose(rhs.grad, rhs_copy.grad, rtol=0.03, atol=1e-5)
+                if lhs is not None:
+                    self.assertAllClose(lhs.grad, lhs_copy.grad, rtol=0.03, atol=1e-5)
+
+            # Determine if we've called CG or not
+            if not cholesky and self.__class__.should_call_cg:
+                self.assertTrue(linear_cg_mock.called)
+            else:
+                self.assertFalse(linear_cg_mock.called)
+
+    def _test_inv_quad_logdet(self, reduce_inv_quad=True, cholesky=False):
+        if not self.__class__.skip_slq_tests:
+            # Forward
+            lazy_tensor = self.create_lazy_tensor()
+            evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+            flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
+
+            vecs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 3, requires_grad=True)
+            vecs_copy = vecs.clone().detach_().requires_grad_(True)
+
+            _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
+            with patch("gpytorch.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
+                with gpytorch.settings.num_trace_samples(128), \
+                        gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), \
+                        gpytorch.settings.cg_tolerance(1e-5):
+
+                    res_inv_quad, res_logdet = lazy_tensor.inv_quad_logdet(
+                        inv_quad_rhs=vecs, logdet=True, reduce_inv_quad=reduce_inv_quad
+                    )
+
+            actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum(-2)
+            if reduce_inv_quad:
+                actual_inv_quad = actual_inv_quad.sum(-1)
+            actual_logdet = torch.cat(
+                [torch.logdet(flattened_evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.batch_shape.numel())]
+            ).view(lazy_tensor.batch_shape)
+
+            self.assertAllClose(res_inv_quad, actual_inv_quad, rtol=0.01, atol=0.01)
+            self.assertAllClose(res_logdet, actual_logdet, rtol=0.2, atol=0.03)
+
+            if not cholesky and self.__class__.should_call_cg:
+                self.assertTrue(linear_cg_mock.called)
+            else:
+                self.assertFalse(linear_cg_mock.called)
 
     def test_add_diag(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -343,19 +330,19 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
         actual = evaluated + torch.eye(evaluated.size(-1)).view(
             *[1 for _ in range(lazy_tensor.dim() - 2)], evaluated.size(-1), evaluated.size(-1)
         ).repeat(*lazy_tensor.batch_shape, 1, 1).mul(1.5)
-        self.assertTrue(approx_equal(res, actual))
+        self.assertAllClose(res, actual)
 
         other_diag = torch.tensor([1.5])
         res = lazy_tensor.add_diag(other_diag).evaluate()
         actual = evaluated + torch.eye(evaluated.size(-1)).view(
             *[1 for _ in range(lazy_tensor.dim() - 2)], evaluated.size(-1), evaluated.size(-1)
         ).repeat(*lazy_tensor.batch_shape, 1, 1).mul(1.5)
-        self.assertTrue(approx_equal(res, actual))
+        self.assertAllClose(res, actual)
 
         other_diag = torch.randn(lazy_tensor.size(-1)).pow(2)
         res = lazy_tensor.add_diag(other_diag).evaluate()
         actual = evaluated + other_diag.diag().repeat(*lazy_tensor.batch_shape, 1, 1)
-        self.assertTrue(approx_equal(res, actual))
+        self.assertAllClose(res, actual)
 
         for sizes in product([1, None], repeat=(lazy_tensor.dim() - 2)):
             batch_shape = [lazy_tensor.batch_shape[i] if size is None else size for i, size in enumerate(sizes)]
@@ -364,7 +351,7 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
             actual = evaluated.clone().detach()
             for i in range(other_diag.size(-1)):
                 actual[..., i, i] = actual[..., i, i] + other_diag[..., i]
-            self.assertTrue(approx_equal(res, actual))
+            self.assertAllClose(res, actual, rtol=1e-2, atol=1e-5)
 
     def test_diag(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -373,236 +360,78 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
         res = lazy_tensor.diag()
         actual = evaluated.diagonal(dim1=-2, dim2=-1)
         actual = actual.view(*lazy_tensor.batch_shape, -1)
-        self.assertEqual(res.size(), lazy_tensor.size()[:-1])
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+        self.assertAllClose(res, actual, rtol=1e-2, atol=1e-5)
 
-    def test_inv_matmul_vec(self):
-        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+    def test_inv_matmul_vector(self, cholesky=False):
+        lazy_tensor = self.create_lazy_tensor()
+        rhs = torch.randn(lazy_tensor.size(-1))
 
         # We skip this test if we're dealing with batch LazyTensors
         # They shouldn't multiply by a vec
         if lazy_tensor.ndimension() > 2:
             return
+        else:
+            return self._test_inv_matmul(rhs)
 
-        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
-        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
-
-        test_vector = torch.randn(lazy_tensor.size(-1), requires_grad=True)
-        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
-        with gpytorch.settings.max_cg_iterations(200):
-            res = lazy_tensor.inv_matmul(test_vector)
-        actual = evaluated.inverse().matmul(test_vector_copy)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-        grad = torch.randn_like(res)
-        res.backward(gradient=grad)
-        actual.backward(gradient=grad)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-
-    def test_inv_matmul_vector_with_left(self):
-        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+    def test_inv_matmul_vector_with_left(self, cholesky=False):
+        lazy_tensor = self.create_lazy_tensor()
+        rhs = torch.randn(lazy_tensor.size(-1))
+        lhs = torch.randn(6, lazy_tensor.size(-1))
 
         # We skip this test if we're dealing with batch LazyTensors
         # They shouldn't multiply by a vec
         if lazy_tensor.ndimension() > 2:
             return
+        else:
+            return self._test_inv_matmul(rhs, lhs=lhs)
 
-        lazy_tensor_copy = lazy_tensor.clone().detach_().requires_grad_(True)
-        evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)
-
-        test_vector = torch.randn(lazy_tensor.size(-1), requires_grad=True)
-        test_left = torch.randn(6, lazy_tensor.size(-1), requires_grad=True)
-        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
-        test_left_copy = test_left.clone().detach().requires_grad_(True)
-        with gpytorch.settings.max_cg_iterations(100):
-            res = lazy_tensor.inv_matmul(test_vector, test_left)
-        actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+    def test_inv_matmul_vector_with_left_cholesky(self):
+        lazy_tensor = self.create_lazy_tensor()
+        rhs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
+        lhs = torch.randn(*lazy_tensor.batch_shape, 6, lazy_tensor.size(-1))
+        return self._test_inv_matmul(rhs, lhs=lhs, cholesky=True)
 
-        grad = torch.randn_like(res)
-        res.backward(gradient=grad)
-        actual.backward(gradient=grad)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-
-    def test_inv_matmul_matrix(self):
-        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)
+    def test_inv_matmul_matrix(self, cholesky=False):
+        lazy_tensor = self.create_lazy_tensor()
+        rhs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
+        return self._test_inv_matmul(rhs, cholesky=cholesky)
 
-        test_vector = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5, requires_grad=True)
-        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
-        with gpytorch.settings.max_cg_iterations(100):
-            res = lazy_tensor.inv_matmul(test_vector)
-        actual = evaluated.inverse().matmul(test_vector_copy)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+    def test_inv_matmul_matrix_cholesky(self):
+        return self.test_inv_matmul_matrix(cholesky=True)
 
-        grad = torch.randn_like(res)
-        res.backward(gradient=grad)
-        actual.backward(gradient=grad)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
+    def test_inv_matmul_matrix_with_left(self):
+        lazy_tensor = self.create_lazy_tensor()
+        rhs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
+        lhs = torch.randn(*lazy_tensor.batch_shape, 3, lazy_tensor.size(-1))
+        return self._test_inv_matmul(rhs, lhs=lhs)
 
     def test_inv_matmul_matrix_broadcast(self):
-        # Right hand size has one more batch dimension
-        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)
-        test_vector = torch.randn(3, *lazy_tensor.batch_shape, lazy_tensor.size(-1), 5, requires_grad=True)
-        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
-        with gpytorch.settings.max_cg_iterations(100):
-            res = lazy_tensor.inv_matmul(test_vector)
-        actual = evaluated.inverse().matmul(test_vector_copy)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+        lazy_tensor = self.create_lazy_tensor()
 
-        grad = torch.randn_like(res)
-        res.backward(gradient=grad)
-        actual.backward(gradient=grad)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
+        # Right hand size has one more batch dimension
+        batch_shape = torch.Size((3, *lazy_tensor.batch_shape))
+        rhs = torch.randn(*batch_shape, lazy_tensor.size(-1), 5)
+        self._test_inv_matmul(rhs)
 
         if lazy_tensor.ndimension() > 2:
             # Right hand size has one fewer batch dimension
-            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)
-            test_vector = torch.randn(*lazy_tensor.batch_shape[1:], lazy_tensor.size(-1), 5)
-            test_vector_copy = test_vector.clone().detach().requires_grad_(True)
-            with gpytorch.settings.max_cg_iterations(100):
-                res = lazy_tensor.inv_matmul(test_vector)
-            actual = evaluated.inverse().matmul(test_vector_copy)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-            grad = torch.randn_like(res)
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-                if arg_copy.grad is not None:
-                    self.assertLess(
-                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                    )
+            batch_shape = torch.Size(lazy_tensor.batch_shape[1:])
+            rhs = torch.randn(*batch_shape, lazy_tensor.size(-1), 5)
+            self._test_inv_matmul(rhs)
 
             # Right hand size has a singleton dimension
-            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)
-            test_vector = torch.randn(*lazy_tensor.batch_shape[:-1], 1, lazy_tensor.size(-1), 5, requires_grad=True)
-            test_vector_copy = test_vector.clone().detach().requires_grad_(True)
-            with gpytorch.settings.max_cg_iterations(100):
-                res = lazy_tensor.inv_matmul(test_vector)
-            actual = evaluated.inverse().matmul(test_vector_copy)
-            self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-            grad = torch.randn_like(res)
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-                if arg_copy.grad is not None:
-                    self.assertLess(
-                        ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                    )
-
-    def test_inv_matmul_matrix_with_left(self):
-        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)
-
-        test_vector = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5, requires_grad=True)
-        test_left = torch.randn(*lazy_tensor.batch_shape, 6, lazy_tensor.size(-1), requires_grad=True)
-        test_vector_copy = test_vector.clone().detach().requires_grad_(True)
-        test_left_copy = test_left.clone().detach().requires_grad_(True)
-        with gpytorch.settings.max_cg_iterations(100):
-            res = lazy_tensor.inv_matmul(test_vector, test_left)
-        actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-        grad = torch.randn_like(res)
-        res.backward(gradient=grad)
-        actual.backward(gradient=grad)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
+            batch_shape = torch.Size((*lazy_tensor.batch_shape[:-1], 1))
+            rhs = torch.randn(*batch_shape, lazy_tensor.size(-1), 5)
+            self._test_inv_matmul(rhs)
 
     def test_inv_quad_logdet(self):
-        if not self.__class__.skip_slq_tests:
-            # Forward
-            lazy_tensor = self.create_lazy_tensor()
-            evaluated = self.evaluate_lazy_tensor(lazy_tensor)
-            flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
-
-            vecs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 3, requires_grad=True)
-            vecs_copy = vecs.clone().detach_().requires_grad_(True)
-
-            with gpytorch.settings.num_trace_samples(128):
-                res_inv_quad, res_logdet = lazy_tensor.inv_quad_logdet(inv_quad_rhs=vecs, logdet=True)
-
-            actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum(-2).sum(-1)
-            actual_logdet = torch.cat(
-                [torch.logdet(flattened_evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.batch_shape.numel())]
-            ).view(lazy_tensor.batch_shape)
-
-            diff_invq = (res_inv_quad - actual_inv_quad).abs() / actual_inv_quad.abs().clamp(1, math.inf)
-            diff_logdet = (res_logdet - actual_logdet).abs() / actual_logdet.abs().clamp(1, math.inf)
-            self.assertLess(diff_invq.max().item(), 0.01)
-            self.assertLess(diff_logdet.max().item(), 0.3)
+        return self._test_inv_quad_logdet(reduce_inv_quad=False, cholesky=False)
 
     def test_inv_quad_logdet_no_reduce(self):
-        if not self.__class__.skip_slq_tests:
-            # Forward
-            lazy_tensor = self.create_lazy_tensor()
-            evaluated = self.evaluate_lazy_tensor(lazy_tensor)
-            flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
+        return self._test_inv_quad_logdet(reduce_inv_quad=True, cholesky=False)
 
-            vecs = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 3, requires_grad=True)
-            vecs_copy = vecs.clone().detach_().requires_grad_(True)
-
-            with gpytorch.settings.num_trace_samples(128):
-                res_inv_quad, res_logdet = lazy_tensor.inv_quad_logdet(
-                    inv_quad_rhs=vecs, logdet=True, reduce_inv_quad=False
-                )
-
-            actual_inv_quad = evaluated.inverse().matmul(vecs_copy).mul(vecs_copy).sum(-2).sum(-1)
-            actual_logdet = torch.cat(
-                [torch.logdet(flattened_evaluated[i]).unsqueeze(0) for i in range(lazy_tensor.batch_shape.numel())]
-            ).view(lazy_tensor.batch_shape)
-
-            diff_invq = (res_inv_quad.sum(-1) - actual_inv_quad).abs() / actual_inv_quad.abs().clamp(1, math.inf)
-            diff_logdet = (res_logdet - actual_logdet).abs() / res_logdet.abs().clamp(1, math.inf)
-            self.assertLess(diff_invq.max().item(), 0.01)
-            self.assertLess(diff_logdet.max().item(), 0.3)
+    def test_inv_quad_logdet_no_reduce_cholesky(self):
+        return self._test_inv_quad_logdet(reduce_inv_quad=True, cholesky=True)
 
     def test_prod(self):
         with gpytorch.settings.fast_computations(covar_root_decomposition=False):
@@ -610,31 +439,29 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
             evaluated = self.evaluate_lazy_tensor(lazy_tensor)
 
             if lazy_tensor.ndimension() > 2:
-                self.assertTrue(
-                    torch.allclose(lazy_tensor.prod(-3).evaluate(), evaluated.prod(-3), atol=1e-2, rtol=1e-2)
-                )
+                self.assertAllClose(lazy_tensor.prod(-3).evaluate(), evaluated.prod(-3), atol=1e-2, rtol=1e-2)
             if lazy_tensor.ndimension() > 3:
-                self.assertTrue(
-                    torch.allclose(lazy_tensor.prod(-4).evaluate(), evaluated.prod(-4), atol=1e-2, rtol=1e-2)
-                )
+                self.assertAllClose(lazy_tensor.prod(-4).evaluate(), evaluated.prod(-4), atol=1e-2, rtol=1e-2)
 
-    def test_root_decomposition(self):
-        # Test with Cholesky
-        lazy_tensor = self.create_lazy_tensor()
-        test_mat = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
-        with gpytorch.settings.max_cholesky_numel(lazy_tensor.matrix_shape.numel() + 1):
-            root_approx = lazy_tensor.root_decomposition()
-            res = root_approx.matmul(test_mat)
-            actual = lazy_tensor.matmul(test_mat)
-            self.assertLess(torch.norm(res - actual) / actual.norm(), 0.1)
-
-        # Test with Lanczos
-        lazy_tensor = self.create_lazy_tensor()
-        with gpytorch.settings.max_cholesky_numel(lazy_tensor.matrix_shape.numel() - 1):
-            root_approx = lazy_tensor.root_decomposition()
-            res = root_approx.matmul(test_mat)
-            actual = lazy_tensor.matmul(test_mat)
-            self.assertLess(torch.norm(res - actual) / actual.norm(), 0.1)
+    def test_root_decomposition(self, cholesky=False):
+        _wrapped_lanczos = MagicMock(wraps=gpytorch.utils.lanczos.lanczos_tridiag)
+        with patch("gpytorch.utils.lanczos.lanczos_tridiag", new=_wrapped_lanczos) as lanczos_mock:
+            lazy_tensor = self.create_lazy_tensor()
+            test_mat = torch.randn(*lazy_tensor.batch_shape, lazy_tensor.size(-1), 5)
+            with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0):
+                root_approx = lazy_tensor.root_decomposition()
+                res = root_approx.matmul(test_mat)
+                actual = lazy_tensor.matmul(test_mat)
+                self.assertAllClose(res, actual, rtol=0.05)
+
+            # Make sure that we're calling the correct function
+            if not cholesky and self.__class__.should_call_lanczos:
+                self.assertTrue(lanczos_mock.called)
+            else:
+                self.assertFalse(lanczos_mock.called)
+
+    def test_root_decomposition_cholesky(self):
+        return self.test_root_decomposition(cholesky=True)
 
     def test_root_inv_decomposition(self):
         lazy_tensor = self.create_lazy_tensor()
@@ -644,7 +471,7 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
 
         res = root_approx.matmul(test_mat)
         actual = lazy_tensor.inv_matmul(test_mat)
-        self.assertLess(torch.norm(res - actual) / actual.norm(), 0.1)
+        self.assertAllClose(res, actual, rtol=0.05, atol=0.02)
 
     def test_sample(self):
         if self.__class__.should_test_sample:
@@ -653,4 +480,4 @@ class LazyTensorTestCase(RectangularLazyTensorTestCase):
 
             samples = lazy_tensor.zero_mean_mvn_samples(50000)
             sample_covar = samples.unsqueeze(-1).matmul(samples.unsqueeze(-2)).mean(0)
-            self.assertLess(((sample_covar - evaluated).abs() / evaluated.abs().clamp(1, math.inf)).max().item(), 3e-1)
+            self.assertAllClose(sample_covar, evaluated, rtol=0.3, atol=0.3)
diff --git a/test/lazy/test_cached_cg_lazy_tensor.py b/test/lazy/test_cached_cg_lazy_tensor.py
index 8e8fae59..45d7899c 100644
--- a/test/lazy/test_cached_cg_lazy_tensor.py
+++ b/test/lazy/test_cached_cg_lazy_tensor.py
@@ -7,155 +7,124 @@ import unittest
 import warnings
 from gpytorch.lazy import CachedCGLazyTensor, NonLazyTensor
 from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
+from unittest.mock import MagicMock, patch
 
 
 class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
     seed = 0
 
-    def create_lazy_tensor(self):
+    def create_lazy_tensor(self, with_solves=False, with_logdet=False):
         mat = torch.randn(5, 6)
         mat = mat.matmul(mat.transpose(-1, -2))
         mat.requires_grad_(True)
 
         lazy_tensor = NonLazyTensor(mat)
         eager_rhs = torch.randn(5, 10).detach()
-        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_logdet tests
-            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
-                lazy_tensor, eager_rhs.detach(), logdet_terms=False
-            )
-            eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
-            solves = [solve.detach(), solve[..., -2:-1].detach()]
+        if with_solves:
+            with gpytorch.settings.num_trace_samples(1000 if with_logdet else 1):  # For inv_quad_logdet tests
+                solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                    lazy_tensor, eager_rhs.detach(), logdet_terms=with_logdet
+                )
+                eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
+                solves = [solve.detach(), solve[..., -2:-1].detach()]
+        else:
+            eager_rhss = [eager_rhs]
+            solves = []
+            probe_vecs = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_norms = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_solves = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            tmats = torch.tensor([], dtype=mat.dtype, device=mat.device)
 
         return CachedCGLazyTensor(lazy_tensor, eager_rhss, solves, probe_vecs, probe_vec_norms, probe_vec_solves, tmats)
 
     def evaluate_lazy_tensor(self, lazy_tensor):
         return lazy_tensor.base_lazy_tensor.tensor
 
-    def test_inv_matmul_vec(self):
-        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)
+    def _test_inv_matmul(self, rhs, lhs=None, cholesky=False):
+        if cholesky:  # These tests don't make sense for CachedCGLazyTensor
+            return
 
-        test_vector = lazy_tensor.eager_rhss[1].squeeze(-1).clone().detach().requires_grad_(True)
-        test_vector_copy = lazy_tensor_copy.eager_rhss[1].squeeze(-1).clone().detach().requires_grad_(True)
-        # Make sure that we get no warning about CG
-        with gpytorch.settings.max_cg_iterations(200), warnings.catch_warnings(record=True) as w:
-            res = lazy_tensor.inv_matmul(test_vector)
-            actual = evaluated.inverse().matmul(test_vector_copy)
-            self.assertEqual(len(w), 0)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-        grad = torch.randn_like(res)
-        # Make sure that we get a warning that CG was run
-        with warnings.catch_warnings(record=True) as w:
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            self.assertEqual(len(w), 1)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-
-    def test_inv_matmul_vector_with_left(self):
-        lazy_tensor = self.create_lazy_tensor().requires_grad_(True)
+        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)
 
-        test_vector = lazy_tensor.eager_rhss[0][..., -1].squeeze(-1).clone().detach().requires_grad_(True)
-        test_vector_copy = lazy_tensor_copy.eager_rhss[0][..., -1].squeeze(-1).clone().detach().requires_grad_(True)
-        test_left = lazy_tensor.eager_rhss[0][..., :-1].t().clone().detach().requires_grad_(True)
-        test_left_copy = lazy_tensor_copy.eager_rhss[0][..., :-1].t().clone().detach().requires_grad_(True)
-        # Make sure that we get no warning about CG
-        with gpytorch.settings.max_cg_iterations(200), warnings.catch_warnings(record=True) as w:
-            res = lazy_tensor.inv_matmul(test_vector, test_left)
-            actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
-            self.assertEqual(len(w), 0)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-        grad = torch.randn_like(res)
-        # Make sure that we get no warning about CG
-        with warnings.catch_warnings(record=True) as w:
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            self.assertEqual(len(w), 0)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-
-    def test_inv_matmul_matrix(self):
-        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)
-
-        test_vector = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
-        test_vector_copy = lazy_tensor_copy.eager_rhss[0].clone().detach().requires_grad_(True)
-        # Make sure that we get no warning about CG
-        with gpytorch.settings.max_cg_iterations(100), warnings.catch_warnings(record=True) as w:
-            res = lazy_tensor.inv_matmul(test_vector)
-            actual = evaluated.inverse().matmul(test_vector_copy)
-            self.assertEqual(len(w), 0)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
-
-        grad = torch.randn_like(res)
-        # Make sure that we get a warning that CG was run
-        with warnings.catch_warnings(record=True) as w:
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            self.assertEqual(len(w), 1)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-
-    def test_inv_matmul_matrix_with_left(self):
-        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)
-
-        test_vector = lazy_tensor.eager_rhss[0][..., 2:].clone().detach().requires_grad_(True)
-        test_vector_copy = lazy_tensor_copy.eager_rhss[0][..., 2:].clone().detach().requires_grad_(True)
-        test_left = lazy_tensor.eager_rhss[0][..., :2].transpose(-1, -2).clone().detach().requires_grad_(True)
-        test_left_copy = lazy_tensor_copy.eager_rhss[0][..., :2].transpose(-1, -2).clone().detach().requires_grad_(True)
-        # Make sure that we get no warning about CG
-        with gpytorch.settings.max_cg_iterations(100), warnings.catch_warnings(record=True) as w:
-            res = lazy_tensor.inv_matmul(test_vector, test_left)
-            actual = test_left_copy @ evaluated.inverse() @ test_vector_copy
-            self.assertEqual(len(w), 0)
-        self.assertLess(((res - actual).abs() / actual.abs().clamp(1, 1e5)).max().item(), 3e-1)
+        # Rather than the supplied rhs and lhs,
+        # we'll replace them with ones that we've precomputed solves for
+        rhs_orig = rhs
+        if rhs_orig.dim() == 1:
+            rhs = lazy_tensor.eager_rhss[0][..., -1].squeeze(-1).clone().detach().requires_grad_(True)
+            # Make sure we're setting this test up correctly
+            self.assertEqual(rhs_orig.shape, rhs.shape)
+        else:
+            if lhs is not None:
+                rhs = lazy_tensor.eager_rhss[0][..., 2:].clone().detach().requires_grad_(True)
+            else:
+                rhs = lazy_tensor.eager_rhss[0].clone().detach().requires_grad_(True)
+            # Make sure we're setting this test up correctly
+            self.assertEqual(rhs_orig.shape[:-1], rhs.shape[:-1])
+
+        lhs = lhs
+        if lhs is not None:
+            lhs_orig = lhs
+            if rhs_orig.dim() == 1:
+                lhs = lazy_tensor.eager_rhss[0][..., :-1].transpose(-1, -2).clone().detach().requires_grad_(True)
+            else:
+                lhs = lazy_tensor.eager_rhss[0][..., :2].transpose(-1, -2).clone().detach().requires_grad_(True)
+
+            # Make sure we're setting this test up correctly
+            self.assertEqual(lhs_orig.shape[:-2], lhs.shape[:-2])
+
+        # Create a test right hand side and left hand side
+        rhs.requires_grad_(True)
+        rhs_copy = rhs.clone().detach().requires_grad_(True)
+        if lhs is not None:
+            lhs.requires_grad_(True)
+            lhs_copy = lhs.clone().detach().requires_grad_(True)
+
+        _wrapped_cg = MagicMock(wraps=gpytorch.utils.linear_cg)
+        with patch("gpytorch.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
+            with gpytorch.settings.max_cholesky_size(math.inf if cholesky else 0), \
+                    gpytorch.settings.cg_tolerance(1e-4):
+                with warnings.catch_warnings(record=True) as w:
+                    # Perform the inv_matmul
+                    if lhs is not None:
+                        res = lazy_tensor.inv_matmul(rhs, lhs)
+                        actual = lhs_copy @ evaluated.inverse() @ rhs_copy
+                    else:
+                        res = lazy_tensor.inv_matmul(rhs)
+                        actual = evaluated.inverse().matmul(rhs_copy)
+                    self.assertAllClose(res, actual, rtol=0.02, atol=1e-5)
+                    self.assertEqual(len(w), 0)
+
+                with warnings.catch_warnings(record=True) as w:
+                    # Perform backward pass
+                    grad = torch.randn_like(res)
+                    res.backward(gradient=grad)
+                    actual.backward(gradient=grad)
+                    for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
+                        if arg_copy.grad is not None:
+                            self.assertAllClose(arg.grad, arg_copy.grad, rtol=0.03, atol=1e-5)
+                    self.assertAllClose(rhs.grad, rhs_copy.grad, rtol=0.03, atol=1e-5)
+                    if lhs is not None:
+                        self.assertAllClose(lhs.grad, lhs_copy.grad, rtol=0.03, atol=1e-5)
+
+                    # Determine if we've called CG or not
+                    # We shouldn't if we supplied a lhs
+                    if lhs is None:
+                        self.assertEqual(len(w), 1)
+                        if not cholesky and self.__class__.should_call_cg:
+                            self.assertTrue(linear_cg_mock.called)
+                    else:
+                        self.assertEqual(len(w), 0)
+                        self.assertFalse(linear_cg_mock.called)
+
+    def test_inv_matmul_vector(self):
+        # Skipping this test because it's not really necessary for CachedCGLazyTensor
+        # We'll only ever be performing inv_matmul against matrices ,r owhen a left hand side is supplied
+        pass
 
-        grad = torch.randn_like(res)
-        # Make sure that we get no warning about CG
-        with warnings.catch_warnings(record=True) as w:
-            res.backward(gradient=grad)
-            actual.backward(gradient=grad)
-            self.assertEqual(len(w), 0)
-        for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
-            if arg_copy.grad is not None:
-                self.assertLess(
-                    ((arg.grad - arg_copy.grad).abs() / arg_copy.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-                )
-        self.assertLess(
-            ((test_left.grad - test_left_copy.grad).abs() / test_left.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
-        self.assertLess(
-            ((test_vector.grad - test_vector_copy.grad).abs() / test_vector.grad.abs().clamp(1, 1e5)).max().item(), 3e-1
-        )
+    def test_inv_matmul_matrix_broadcast(self):
+        pass
 
     def test_inv_quad_logdet(self):
         pass
@@ -173,35 +142,13 @@ class TestCachedCGLazyTensorNoLogdet(LazyTensorTestCase, unittest.TestCase):
         actual = lazy_tensor.inv_matmul(test_mat)
         self.assertLess(torch.norm(res - actual) / actual.norm(), 0.1)
 
-    def test_inv_matmul_matrix_broadcast(self):
-        pass
-
 
 class TestCachedCGLazyTensor(TestCachedCGLazyTensorNoLogdet):
     seed = 0
 
-    def create_lazy_tensor(self):
-        mat = torch.randn(5, 6)
-        mat = mat.matmul(mat.transpose(-1, -2))
-        mat.requires_grad_(True)
-
-        lazy_tensor = NonLazyTensor(mat)
-        eager_rhs = torch.randn(5, 10).detach()
-        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_logdet tests
-            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
-                lazy_tensor, eager_rhs.detach()
-            )
-            eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
-            solves = [solve.detach(), solve[..., -2:-1].detach()]
-
-        return CachedCGLazyTensor(lazy_tensor, eager_rhss, solves, probe_vecs, probe_vec_norms, probe_vec_solves, tmats)
-
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        return lazy_tensor.base_lazy_tensor.tensor
-
     def test_inv_quad_logdet(self):
         # Forward
-        lazy_tensor = self.create_lazy_tensor()
+        lazy_tensor = self.create_lazy_tensor(with_solves=True, with_logdet=True)
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
         flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
 
@@ -224,7 +171,7 @@ class TestCachedCGLazyTensor(TestCachedCGLazyTensorNoLogdet):
 
     def test_inv_quad_logdet_no_reduce(self):
         # Forward
-        lazy_tensor = self.create_lazy_tensor()
+        lazy_tensor = self.create_lazy_tensor(with_solves=True, with_logdet=True)
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
         flattened_evaluated = evaluated.view(-1, *lazy_tensor.matrix_shape)
 
@@ -251,60 +198,62 @@ class TestCachedCGLazyTensor(TestCachedCGLazyTensorNoLogdet):
 class TestCachedCGLazyTensorNoLogdetBatch(TestCachedCGLazyTensorNoLogdet):
     seed = 0
 
-    def create_lazy_tensor(self):
+    def create_lazy_tensor(self, with_solves=False, with_logdet=False):
         mat = torch.randn(3, 5, 6)
         mat = mat.matmul(mat.transpose(-1, -2))
         mat.requires_grad_(True)
 
-        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_logdet tests
-            lazy_tensor = NonLazyTensor(mat)
-            eager_rhs = torch.randn(3, 5, 10).detach()
-            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
-                lazy_tensor, eager_rhs.detach(), logdet_terms=False
-            )
+        lazy_tensor = NonLazyTensor(mat)
+        eager_rhs = torch.randn(3, 5, 10).detach()
+        if with_solves:
+            with gpytorch.settings.num_trace_samples(1000 if with_logdet else 1):  # For inv_quad_logdet tests
+                solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                    lazy_tensor, eager_rhs.detach(), logdet_terms=with_logdet
+                )
+                eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
+                solves = [solve.detach(), solve[..., -2:-1].detach()]
+        else:
+            eager_rhss = [eager_rhs]
+            solves = []
+            probe_vecs = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_norms = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_solves = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            tmats = torch.tensor([], dtype=mat.dtype, device=mat.device)
 
         return CachedCGLazyTensor(
-            lazy_tensor, [eager_rhs], [solve], probe_vecs, probe_vec_norms, probe_vec_solves, tmats
+            lazy_tensor, eager_rhss, solves, probe_vecs, probe_vec_norms, probe_vec_solves, tmats
         )
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        return lazy_tensor.base_lazy_tensor.tensor
-
-    def test_inv_matmul_vec(self):
-        pass
-
-    def test_inv_matmul_vector_with_left(self):
-        pass
-
 
 class TestCachedCGLazyTensorBatch(TestCachedCGLazyTensor):
     seed = 0
 
-    def create_lazy_tensor(self):
+    def create_lazy_tensor(self, with_solves=False, with_logdet=False):
         mat = torch.randn(3, 5, 6)
         mat = mat.matmul(mat.transpose(-1, -2))
         mat.requires_grad_(True)
 
-        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_logdet tests
-            lazy_tensor = NonLazyTensor(mat)
-            eager_rhs = torch.randn(3, 5, 10).detach()
-            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
-                lazy_tensor, eager_rhs.detach()
-            )
+        lazy_tensor = NonLazyTensor(mat)
+        eager_rhs = torch.randn(3, 5, 10).detach()
+        if with_solves:
+            with gpytorch.settings.num_trace_samples(1000 if with_logdet else 1):  # For inv_quad_logdet tests
+                solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                    lazy_tensor, eager_rhs.detach(), logdet_terms=with_logdet
+                )
+                eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
+                solves = [solve.detach(), solve[..., -2:-1].detach()]
+        else:
+            eager_rhss = [eager_rhs]
+            solves = []
+            probe_vecs = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_norms = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_solves = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            tmats = torch.tensor([], dtype=mat.dtype, device=mat.device)
 
         return CachedCGLazyTensor(
-            lazy_tensor, [eager_rhs], [solve], probe_vecs, probe_vec_norms, probe_vec_solves, tmats
+            lazy_tensor, eager_rhss, solves, probe_vecs, probe_vec_norms, probe_vec_solves, tmats
         )
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        return lazy_tensor.base_lazy_tensor.tensor
-
-    def test_inv_matmul_vec(self):
-        pass
-
-    def test_inv_matmul_vector_with_left(self):
-        pass
-
 
 class TestCachedCGLazyTensorMultiBatch(TestCachedCGLazyTensor):
     seed = 0
@@ -312,27 +261,28 @@ class TestCachedCGLazyTensorMultiBatch(TestCachedCGLazyTensor):
     should_test_sample = False
     skip_slq_tests = True
 
-    def create_lazy_tensor(self):
+    def create_lazy_tensor(self, with_solves=False, with_logdet=False):
         mat = torch.randn(2, 3, 5, 6)
         mat = mat.matmul(mat.transpose(-1, -2))
         mat.requires_grad_(True)
 
-        with gpytorch.settings.num_trace_samples(1000):  # For inv_quad_logdet tests
-            lazy_tensor = NonLazyTensor(mat)
-            eager_rhs = torch.randn(2, 3, 5, 10).detach()
-            solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
-                lazy_tensor, eager_rhs.detach()
-            )
+        lazy_tensor = NonLazyTensor(mat)
+        eager_rhs = torch.randn(2, 3, 5, 10).detach()
+        if with_solves:
+            with gpytorch.settings.num_trace_samples(1000 if with_logdet else 1):  # For inv_quad_logdet tests
+                solve, probe_vecs, probe_vec_norms, probe_vec_solves, tmats = CachedCGLazyTensor.precompute_terms(
+                    lazy_tensor, eager_rhs.detach(), logdet_terms=with_logdet
+                )
+                eager_rhss = [eager_rhs.detach(), eager_rhs[..., -2:-1].detach()]
+                solves = [solve.detach(), solve[..., -2:-1].detach()]
+        else:
+            eager_rhss = [eager_rhs]
+            solves = []
+            probe_vecs = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_norms = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            probe_vec_solves = torch.tensor([], dtype=mat.dtype, device=mat.device)
+            tmats = torch.tensor([], dtype=mat.dtype, device=mat.device)
 
         return CachedCGLazyTensor(
-            lazy_tensor, [eager_rhs], [solve], probe_vecs, probe_vec_norms, probe_vec_solves, tmats
+            lazy_tensor, eager_rhss, solves, probe_vecs, probe_vec_norms, probe_vec_solves, tmats
         )
-
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        return lazy_tensor.base_lazy_tensor.tensor
-
-    def test_inv_matmul_vec(self):
-        pass
-
-    def test_inv_matmul_vector_with_left(self):
-        pass
diff --git a/test/lazy/test_chol_lazy_tensor.py b/test/lazy/test_chol_lazy_tensor.py
index 05f57460..1d715c40 100644
--- a/test/lazy/test_chol_lazy_tensor.py
+++ b/test/lazy/test_chol_lazy_tensor.py
@@ -9,6 +9,8 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 class TestCholLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
+    should_call_cg = False
+    should_call_lanczos = False
 
     def create_lazy_tensor(self):
         chol = torch.tensor(
@@ -22,18 +24,9 @@ class TestCholLazyTensor(LazyTensorTestCase, unittest.TestCase):
         chol = lazy_tensor.root.evaluate()
         return chol.matmul(chol.transpose(-1, -2))
 
-    def test_inv_matmul_vec(self):
-        # We're skipping this test, since backward passes aren't defined in torch for this
-        pass
 
-    def test_inv_matmul_matrix(self):
-        # We're skipping this test, since backward passes aren't defined in torch for this
-        pass
-
-
-class TestCholLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
+class TestCholLazyTensorBatch(TestCholLazyTensor):
     seed = 0
-    should_test_sample = True
 
     def create_lazy_tensor(self):
         chol = torch.tensor(
@@ -47,13 +40,8 @@ class TestCholLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
         chol.requires_grad_(True)
         return CholLazyTensor(chol)
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        chol = lazy_tensor.root.evaluate()
-        res = chol.matmul(chol.transpose(-1, -2))
-        return res
 
-
-class TestCholLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
+class TestCholLazyTensorMultiBatch(TestCholLazyTensor):
     seed = 0
     # Because these LTs are large, we'll skil the big tests
     should_test_sample = False
@@ -74,11 +62,6 @@ class TestCholLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
         chol.requires_grad_(True)
         return CholLazyTensor(chol)
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        chol = lazy_tensor.root.evaluate()
-        res = chol.matmul(chol.transpose(-1, -2))
-        return res
-
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/lazy/test_diag_lazy_tensor.py b/test/lazy/test_diag_lazy_tensor.py
index e125d74d..61a6d168 100644
--- a/test/lazy/test_diag_lazy_tensor.py
+++ b/test/lazy/test_diag_lazy_tensor.py
@@ -9,6 +9,8 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 class TestDiagLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
+    should_call_cg = False
+    should_call_lanczos = False
 
     def create_lazy_tensor(self):
         diag = torch.tensor([1.0, 2.0, 4.0, 2.0, 3.0], requires_grad=True)
@@ -19,9 +21,8 @@ class TestDiagLazyTensor(LazyTensorTestCase, unittest.TestCase):
         return diag.diag()
 
 
-class TestDiagLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
+class TestDiagLazyTensorBatch(TestDiagLazyTensor):
     seed = 0
-    should_test_sample = True
 
     def create_lazy_tensor(self):
         diag = torch.tensor(
@@ -34,7 +35,7 @@ class TestDiagLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
         return torch.cat([diag[i].diag().unsqueeze(0) for i in range(3)])
 
 
-class TestDiagLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
+class TestDiagLazyTensorMultiBatch(TestDiagLazyTensor):
     seed = 0
     # Because these LTs are large, we'll skil the big tests
     should_test_sample = True
diff --git a/test/lazy/test_kronecker_product_lazy_tensor.py b/test/lazy/test_kronecker_product_lazy_tensor.py
index 474cf4ff..3f63d267 100644
--- a/test/lazy/test_kronecker_product_lazy_tensor.py
+++ b/test/lazy/test_kronecker_product_lazy_tensor.py
@@ -17,6 +17,8 @@ def kron(a, b):
 
 
 class TestKroneckerProductLazyTensor(LazyTensorTestCase, unittest.TestCase):
+    seed = 0
+
     def create_lazy_tensor(self):
         a = torch.tensor([[4, 0, 2], [0, 3, -1], [2, -1, 3]], dtype=torch.float)
         b = torch.tensor([[2, 1], [1, 2]], dtype=torch.float)
@@ -33,7 +35,9 @@ class TestKroneckerProductLazyTensor(LazyTensorTestCase, unittest.TestCase):
         return res
 
 
-class TestKroneckerProductLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
+class TestKroneckerProductLazyTensorBatch(TestKroneckerProductLazyTensor):
+    seed = 0
+
     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)
@@ -44,13 +48,10 @@ class TestKroneckerProductLazyTensorBatch(LazyTensorTestCase, unittest.TestCase)
         kp_lazy_tensor = KroneckerProductLazyTensor(NonLazyTensor(a), NonLazyTensor(b), NonLazyTensor(c))
         return kp_lazy_tensor
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        res = kron(lazy_tensor.lazy_tensors[0].tensor, lazy_tensor.lazy_tensors[1].tensor)
-        res = kron(res, lazy_tensor.lazy_tensors[2].tensor)
-        return res
-
 
 class TestKroneckerProductLazyTensorRectangular(RectangularLazyTensorTestCase, unittest.TestCase):
+    seed = 0
+
     def create_lazy_tensor(self):
         a = torch.randn(2, 3, requires_grad=True)
         b = torch.randn(5, 2, requires_grad=True)
@@ -64,7 +65,9 @@ class TestKroneckerProductLazyTensorRectangular(RectangularLazyTensorTestCase, u
         return res
 
 
-class TestKroneckerProductLazyTensorRectangularMultiBatch(RectangularLazyTensorTestCase, unittest.TestCase):
+class TestKroneckerProductLazyTensorRectangularMultiBatch(TestKroneckerProductLazyTensorRectangular):
+    seed = 0
+
     def create_lazy_tensor(self):
         a = torch.randn(3, 4, 2, 3, requires_grad=True)
         b = torch.randn(3, 4, 5, 2, requires_grad=True)
@@ -72,11 +75,6 @@ class TestKroneckerProductLazyTensorRectangularMultiBatch(RectangularLazyTensorT
         kp_lazy_tensor = KroneckerProductLazyTensor(NonLazyTensor(a), NonLazyTensor(b), NonLazyTensor(c))
         return kp_lazy_tensor
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        res = kron(lazy_tensor.lazy_tensors[0].tensor, lazy_tensor.lazy_tensors[1].tensor)
-        res = kron(res, lazy_tensor.lazy_tensors[2].tensor)
-        return res
-
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/lazy/test_mul_lazy_tensor.py b/test/lazy/test_mul_lazy_tensor.py
index acf7cc66..32cd521d 100644
--- a/test/lazy/test_mul_lazy_tensor.py
+++ b/test/lazy/test_mul_lazy_tensor.py
@@ -8,7 +8,6 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 
 def make_random_mat(size, rank, batch_shape=torch.Size(())):
     res = torch.randn(*batch_shape, size, rank)
-    res.requires_grad_()
     return res
 
 
@@ -16,8 +15,8 @@ class TestMulLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 10
 
     def create_lazy_tensor(self):
-        mat1 = make_random_mat(6, 3)
-        mat2 = make_random_mat(6, 3)
+        mat1 = make_random_mat(6, 6)
+        mat2 = make_random_mat(6, 6)
         res = RootLazyTensor(mat1) * RootLazyTensor(mat2)
         return res.add_diag(torch.tensor(2.0))
 
@@ -35,8 +34,8 @@ class TestMulLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
     seed = 2
 
     def create_lazy_tensor(self):
-        mat1 = make_random_mat(6, rank=5, batch_shape=torch.Size((2,)))
-        mat2 = make_random_mat(6, rank=5, batch_shape=torch.Size((2,)))
+        mat1 = make_random_mat(6, rank=6, batch_shape=torch.Size((2,)))
+        mat2 = make_random_mat(6, rank=6, batch_shape=torch.Size((2,)))
         res = RootLazyTensor(mat1) * RootLazyTensor(mat2)
         return res.add_diag(torch.tensor(2.0))
 
@@ -55,8 +54,8 @@ class TestMulLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
     skip_slq_tests = True
 
     def create_lazy_tensor(self):
-        mat1 = make_random_mat(6, rank=5, batch_shape=torch.Size((2, 3,)))
-        mat2 = make_random_mat(6, rank=5, batch_shape=torch.Size((2, 3,)))
+        mat1 = make_random_mat(6, rank=6, batch_shape=torch.Size((2, 3,)))
+        mat2 = make_random_mat(6, rank=6, batch_shape=torch.Size((2, 3,)))
         res = RootLazyTensor(mat1) * RootLazyTensor(mat2)
         return res.add_diag(torch.tensor(0.5))
 
diff --git a/test/lazy/test_root_lazy_tensor.py b/test/lazy/test_root_lazy_tensor.py
index a223ea0f..74348bef 100644
--- a/test/lazy/test_root_lazy_tensor.py
+++ b/test/lazy/test_root_lazy_tensor.py
@@ -9,6 +9,7 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase
 class TestRootLazyTensor(LazyTensorTestCase, unittest.TestCase):
     seed = 0
     should_test_sample = True
+    should_call_lanczos = False
 
     def create_lazy_tensor(self):
         root = torch.randn(3, 5, requires_grad=True)
@@ -20,9 +21,8 @@ class TestRootLazyTensor(LazyTensorTestCase, unittest.TestCase):
         return res
 
 
-class TestRootLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
+class TestRootLazyTensorBatch(TestRootLazyTensor):
     seed = 1
-    should_test_sample = True
 
     def create_lazy_tensor(self):
         root = torch.randn(3, 5, 5)
@@ -30,14 +30,9 @@ class TestRootLazyTensorBatch(LazyTensorTestCase, unittest.TestCase):
         root.requires_grad_(True)
         return RootLazyTensor(root)
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        root = lazy_tensor.root.tensor
-        res = root.matmul(root.transpose(-1, -2))
-        return res
-
 
-class TestRootLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
-    seed = 0
+class TestRootLazyTensorMultiBatch(TestRootLazyTensor):
+    seed = 1
     # Because these LTs are large, we'll skil the big tests
     should_test_sample = False
     skip_slq_tests = True
@@ -47,11 +42,6 @@ class TestRootLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
         root.requires_grad_(True)
         return RootLazyTensor(root)
 
-    def evaluate_lazy_tensor(self, lazy_tensor):
-        root = lazy_tensor.root.tensor
-        res = root.matmul(root.transpose(-1, -2))
-        return res
-
 
 if __name__ == "__main__":
     unittest.main()