Repository: cornellius-gp/gpytorch
Pull request #427 (2018-12-12 14:20:45)

Correct order of Kronecker Product
`KroneckerProductLazyTensor(a, b, c)` was computing `c \kron b \kron a`, not `a \kron b \kron c`. This PR fixes that.


Patch:
diff --git a/docs/source/functions.rst b/docs/source/functions.rst
index 037851d1..9cf3fd01 100644
--- a/docs/source/functions.rst
+++ b/docs/source/functions.rst
@@ -18,6 +18,10 @@ Functions
 
 .. autofunction:: dsmm
 
+.. autofunction:: exact_predictive_mean
+
+.. autofunction:: exact_predictive_covar
+
 .. autofunction:: inv_matmul
 
 .. autofunction:: inv_quad
diff --git a/gpytorch/__init__.py b/gpytorch/__init__.py
index 3da541b6..148013c4 100644
--- a/gpytorch/__init__.py
+++ b/gpytorch/__init__.py
@@ -20,6 +20,8 @@ from .functions import (
     add_diag,
     add_jitter,
     dsmm,
+    exact_predictive_covar,
+    exact_predictive_mean,
     inv_matmul,
     inv_quad,
     inv_quad_log_det,
@@ -53,6 +55,8 @@ __all__ = [
     "add_diag",
     "add_jitter",
     "dsmm",
+    "exact_predictive_mean",
+    "exact_predictive_covar",
     "inv_matmul",
     "inv_quad",
     "inv_quad_log_det",
diff --git a/gpytorch/functions/__init__.py b/gpytorch/functions/__init__.py
index 6051899a..ad62644d 100644
--- a/gpytorch/functions/__init__.py
+++ b/gpytorch/functions/__init__.py
@@ -63,6 +63,71 @@ def dsmm(sparse_mat, dense_mat):
     return DSMM(sparse_mat)(dense_mat)
 
 
+def exact_predictive_mean(
+    full_covar,
+    full_mean,
+    train_inputs,
+    train_labels,
+    num_train,
+    likelihood,
+    precomputed_cache=None,
+    non_batch_train=False,
+):
+    """
+    Computes the posterior predictive mean of a GP
+
+    Args:
+        - full_covar ( (n+t) x (n+t) ) - the block prior covariance matrix of training and testing points
+            [ K_XX, K_XX*; K_X*X, K_X*X* ]
+        - full_mean (n + t) - the training and test prior means, stacked on top of each other
+        - train_inputs (:obj:`torch.tensor`) -  The training data inputs
+        - train_labels (n) - the training labels minus the training prior mean
+        - noise (1) - the observed noise (from the likelihood)
+        - precomputed_cache - speeds up subsequent computations (default: None)
+
+    Returns:
+        - (t) - the predictive posterior mean of the test points
+    """
+    if not num_train:
+        return full_mean, None
+
+    if not hasattr(full_covar, "exact_predictive_mean"):
+        from ..lazy.non_lazy_tensor import NonLazyTensor
+
+        full_covar = NonLazyTensor(full_covar)
+    return full_covar.exact_predictive_mean(
+        full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache, non_batch_train
+    )
+
+
+def exact_predictive_covar(
+    full_covar, train_inputs, num_train, likelihood, precomputed_cache=None, non_batch_train=False
+):
+    """
+    Computes the posterior predictive covariance of a GP
+
+    Args:
+        - full_covar ( (n+t) x (n+t) ) - the block prior covariance matrix of training and testing points
+            [ K_XX, K_XX*; K_X*X, K_X*X* ]
+        - train_inputs TODO
+        - num_train (int) - how many training points are there in the full covariance matrix
+        - noise (1) - the observed noise (from the likelihood)
+        - precomputed_cache - speeds up subsequent computations (default: None)
+
+    Returns:
+        - LazyTensor (t x t) - the predictive posterior covariance of the test points
+    """
+    if not num_train:
+        return full_covar, None
+
+    if not hasattr(full_covar, "exact_predictive_covar"):
+        from ..lazy.non_lazy_tensor import NonLazyTensor
+
+        full_covar = NonLazyTensor(full_covar)
+
+    return full_covar.exact_predictive_covar(train_inputs, num_train, likelihood, precomputed_cache, non_batch_train)
+
+
 def log_normal_cdf(x):
     """
     Computes the element-wise log standard normal CDF of an input tensor x.
@@ -193,6 +258,8 @@ def root_inv_decomposition(mat, initial_vectors=None, test_vectors=None):
 __all__ = [
     "add_diag",
     "dsmm",
+    "exact_predictive_mean",
+    "exact_predictive_covar",
     "inv_matmul",
     "inv_quad",
     "inv_quad_log_det",
diff --git a/gpytorch/kernels/grid_interpolation_kernel.py b/gpytorch/kernels/grid_interpolation_kernel.py
index bb24f076..bcb5d672 100644
--- a/gpytorch/kernels/grid_interpolation_kernel.py
+++ b/gpytorch/kernels/grid_interpolation_kernel.py
@@ -113,6 +113,10 @@ class GridInterpolationKernel(GridKernel):
             for bound, spacing in zip(self.grid_bounds, grid_spacings)
         )
 
+    @property
+    def has_custom_exact_predictions(self):
+        return True
+
     def _compute_grid(self, inputs, batch_dims):
         batch_size, n_data, n_dimensions = inputs.size()
         if batch_dims == (0, 2):
diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index 51c9ed4c..aaa7199b 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -88,7 +88,7 @@ class GridKernel(Kernel):
                     covars = [full_covar[i : i + 1] for i in range(n_dim)]
 
             if len(covars) > 1:
-                covar = KroneckerProductLazyTensor(*covars)
+                covar = KroneckerProductLazyTensor(*covars[::-1])
             else:
                 covar = covars[0]
 
diff --git a/gpytorch/kernels/inducing_point_kernel.py b/gpytorch/kernels/inducing_point_kernel.py
index 35dc4144..0312e272 100644
--- a/gpytorch/kernels/inducing_point_kernel.py
+++ b/gpytorch/kernels/inducing_point_kernel.py
@@ -61,7 +61,7 @@ class InducingPointKernel(Kernel):
 
             # Diagonal correction for predictive posterior
             correction = (self.base_kernel(x1, x2).diag() - covar.diag()).clamp(0, math.inf)
-            covar = PsdSumLazyTensor(DiagLazyTensor(correction), covar)
+            covar = PsdSumLazyTensor(DiagLazyTensor(correction) + covar)
         else:
             k_ux2 = self.base_kernel(x2, self.inducing_points).evaluate()
             covar = MatmulLazyTensor(
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 7d770d25..b85ed3c3 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -143,6 +143,10 @@ class Kernel(Module):
             value = torch.tensor(value)
         self.initialize(raw_lengthscale=self._inv_param_transform(value))
 
+    @property
+    def has_custom_exact_predictions(self):
+        return False
+
     def size(self, x1, x2):
         non_batch_size = (x1.size(-2), x2.size(-2))
         if x1.ndimension() == 3:
diff --git a/gpytorch/kernels/multitask_kernel.py b/gpytorch/kernels/multitask_kernel.py
index 758a2b86..6beee9ed 100644
--- a/gpytorch/kernels/multitask_kernel.py
+++ b/gpytorch/kernels/multitask_kernel.py
@@ -48,7 +48,7 @@ class MultitaskKernel(Kernel):
         covar_x = self.data_covar_module.forward(x1, x2, **params)
         if not isinstance(covar_x, LazyTensor):
             covar_x = NonLazyTensor(covar_x)
-        res = KroneckerProductLazyTensor(covar_i, covar_x)
+        res = KroneckerProductLazyTensor(covar_x, covar_i)
         return res.diag() if diag else res
 
     def size(self, x1, x2):
diff --git a/gpytorch/lazy/constant_mul_lazy_tensor.py b/gpytorch/lazy/constant_mul_lazy_tensor.py
index 44fc10ec..16029dcb 100644
--- a/gpytorch/lazy/constant_mul_lazy_tensor.py
+++ b/gpytorch/lazy/constant_mul_lazy_tensor.py
@@ -122,7 +122,6 @@ class ConstantMulLazyTensor(LazyTensor):
         return ConstantMulLazyTensor(self.base_lazy_tensor._transpose_nonbatch(), self.constant)
 
     @property
-    @cached
     def constant(self):
         # Make sure that the constant can be expanded to the appropriate size
         try:
diff --git a/gpytorch/lazy/interpolated_lazy_tensor.py b/gpytorch/lazy/interpolated_lazy_tensor.py
index 058bbfce..1a8a8276 100644
--- a/gpytorch/lazy/interpolated_lazy_tensor.py
+++ b/gpytorch/lazy/interpolated_lazy_tensor.py
@@ -8,6 +8,7 @@ from .root_lazy_tensor import RootLazyTensor
 from ..utils import sparse
 from ..utils.interpolation import left_interp, left_t_interp
 from ..utils.sparse import bdsmm
+from .. import beta_features
 
 
 class InterpolatedLazyTensor(LazyTensor):
@@ -339,6 +340,172 @@ class InterpolatedLazyTensor(LazyTensor):
                 res = res.view(batch_size, n_data, -1).sum(-1)
             return res
 
+    def exact_predictive_mean(
+        self,
+        full_mean,
+        train_inputs,
+        train_labels,
+        num_train,
+        likelihood,
+        precomputed_cache=None,
+        non_batch_train=False,
+    ):
+        from ..distributions import MultivariateNormal
+
+        if precomputed_cache is None:
+            train_interp_indices = self.left_interp_indices.narrow(-2, 0, num_train)
+            train_interp_values = self.left_interp_values.narrow(-2, 0, num_train)
+
+            # Get inverse root
+            train_train_covar = self.__class__(
+                self.base_lazy_tensor,
+                train_interp_indices,
+                train_interp_values,
+                train_interp_indices,
+                train_interp_values,
+            )
+
+            train_mean = full_mean.narrow(-1, 0, train_train_covar.size(-1))
+
+            mvn = likelihood(MultivariateNormal(train_mean, train_train_covar), train_inputs)
+            train_mean, train_train_covar = mvn.mean, mvn.lazy_covariance_matrix
+
+            train_train_covar_inv_labels = train_train_covar.inv_matmul((train_labels - train_mean).unsqueeze(-1))
+
+            # New root factor
+            base_size = self.base_lazy_tensor.size(-1)
+            precomputed_cache = self.base_lazy_tensor.matmul(
+                left_t_interp(train_interp_indices, train_interp_values, train_train_covar_inv_labels, base_size)
+            )
+
+            # Prevent backprop through this variable
+            precomputed_cache = precomputed_cache.detach()
+
+        # Compute the exact predictive posterior
+        n_test = self.size(-1) - num_train
+        test_mean = full_mean.narrow(-1, num_train, n_test)
+        test_interp_indices = self.left_interp_indices.narrow(-2, num_train, n_test)
+        test_interp_values = self.left_interp_values.narrow(-2, num_train, n_test)
+        res = left_interp(test_interp_indices, test_interp_values, precomputed_cache).squeeze(-1) + test_mean
+        return res, precomputed_cache
+
+    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
+        train_interp_indices = test_train_covar.right_interp_indices
+        train_interp_values = test_train_covar.right_interp_values
+        base_lazy_tensor = test_train_covar.base_lazy_tensor
+        base_size = base_lazy_tensor.size(-1)
+        res = base_lazy_tensor.matmul(
+            left_t_interp(train_interp_indices, train_interp_values, train_train_covar_inv_root, base_size)
+        )
+        return res
+
+    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
+        # Here the precomputed cache represents K_UU W S,
+        # where S S^T = (K_XX + sigma^2 I)^-1
+        test_interp_indices = test_train_covar.left_interp_indices
+        test_interp_values = test_train_covar.left_interp_values
+        res = left_interp(test_interp_indices, test_interp_values, precomputed_cache)
+        return res
+
+    def exact_predictive_covar(
+        self, train_inputs, num_train, likelihood, precomputed_cache=None, non_batch_train=False
+    ):
+        from ..distributions import MultivariateNormal
+
+        if not beta_features.fast_pred_var.on() and not beta_features.fast_pred_samples.on():
+            return super(InterpolatedLazyTensor, self).exact_predictive_covar(
+                train_inputs, num_train, likelihood, precomputed_cache
+            )
+
+        n_test = self.size(-2) - num_train
+        train_interp_indices = self.left_interp_indices.narrow(-2, 0, num_train)
+        train_interp_values = self.left_interp_values.narrow(-2, 0, num_train)
+        test_interp_indices = self.left_interp_indices.narrow(-2, num_train, n_test)
+        test_interp_values = self.left_interp_values.narrow(-2, num_train, n_test)
+        test_train_covar = self.__class__(
+            self.base_lazy_tensor, test_interp_indices, test_interp_values, train_interp_indices, train_interp_values
+        )
+
+        if (
+            precomputed_cache is None
+            or (beta_features.fast_pred_samples.on() and precomputed_cache[0] is None)
+            or (not beta_features.fast_pred_samples.on() and precomputed_cache[1] is None)
+        ):
+            # Get inverse root
+            train_train_covar = self.__class__(
+                self.base_lazy_tensor,
+                train_interp_indices,
+                train_interp_values,
+                train_interp_indices,
+                train_interp_values,
+            )
+
+            grv = MultivariateNormal(torch.zeros(1), train_train_covar)
+            train_train_covar = likelihood(grv, train_inputs).lazy_covariance_matrix
+
+            # Get probe vectors for inverse root
+            num_probe_vectors = beta_features.fast_pred_var.num_probe_vectors()
+            batch_size = train_interp_indices.size(0)
+            n_inducing = self.base_lazy_tensor.size(-1)
+            vector_indices = torch.randperm(n_inducing).type_as(train_interp_indices)
+            probe_vector_indices = vector_indices[:num_probe_vectors]
+            test_vector_indices = vector_indices[num_probe_vectors : 2 * num_probe_vectors]
+
+            probe_interp_indices = probe_vector_indices.unsqueeze(1)
+            probe_test_interp_indices = test_vector_indices.unsqueeze(1)
+            probe_interp_values = torch.ones(num_probe_vectors, 1, dtype=self.dtype, device=self.device)
+            if train_interp_indices.ndimension() == 3:
+                probe_interp_indices = probe_interp_indices.unsqueeze(0).expand(batch_size, num_probe_vectors, 1)
+                probe_test_interp_indices = probe_test_interp_indices.unsqueeze(0)
+                probe_test_interp_indices = probe_test_interp_indices.expand(batch_size, num_probe_vectors, 1)
+                probe_interp_values = probe_interp_values.unsqueeze(0).expand(batch_size, num_probe_vectors, 1)
+
+            probe_vectors = InterpolatedLazyTensor(
+                self.base_lazy_tensor,
+                train_interp_indices,
+                train_interp_values,
+                probe_interp_indices,
+                probe_interp_values,
+            ).evaluate()
+            test_vectors = InterpolatedLazyTensor(
+                self.base_lazy_tensor,
+                train_interp_indices,
+                train_interp_values,
+                probe_test_interp_indices,
+                probe_interp_values,
+            ).evaluate()
+
+            # Get inverse root
+            train_train_covar_inv_root = train_train_covar.root_inv_decomposition(probe_vectors, test_vectors).root
+            train_train_covar_inv_root = train_train_covar_inv_root.evaluate()
+
+            # New root factor
+            root = self._exact_predictive_covar_inv_quad_form_cache(train_train_covar_inv_root, test_train_covar)
+
+            # Precomputed factor
+            if beta_features.fast_pred_samples.on():
+                inside = self.base_lazy_tensor + RootLazyTensor(root).mul(-1)
+                inside_root = inside.root_decomposition().root.evaluate()
+                # Prevent backprop through this variable
+                inside_root = inside_root.detach()
+                precomputed_cache = inside_root, None
+            else:
+                # Prevent backprop through this variable
+                root = root.detach()
+                precomputed_cache = None, root
+
+        # Compute the exact predictive posterior
+        if beta_features.fast_pred_samples.on():
+            res = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache[0], test_train_covar)
+            res = RootLazyTensor(res)
+        else:
+            test_test_prior_covar = InterpolatedLazyTensor(
+                self.base_lazy_tensor, test_interp_indices, test_interp_values, test_interp_indices, test_interp_values
+            )
+            root = left_interp(test_interp_indices, test_interp_values, precomputed_cache[1])
+            res = test_test_prior_covar + RootLazyTensor(root).mul(-1)
+        return res, precomputed_cache
+
     def matmul(self, tensor):
         # We're using a custom matmul here, because it is significantly faster than
         # what we get from the function factory.
diff --git a/gpytorch/lazy/kronecker_product_lazy_tensor.py b/gpytorch/lazy/kronecker_product_lazy_tensor.py
index c3ea8c9a..95870012 100644
--- a/gpytorch/lazy/kronecker_product_lazy_tensor.py
+++ b/gpytorch/lazy/kronecker_product_lazy_tensor.py
@@ -12,42 +12,23 @@ def _prod(iterable):
 
 def _matmul(lazy_tensors, rhs):
     res = rhs.contiguous()
-    is_batch = res.ndimension() == 3
     num_cols = rhs.size(-1)
-    for lazy_tensor in list(lazy_tensors)[::-1]:
-        if is_batch:
-            num_batch = res.size(0)
-            res = res.view(num_batch, lazy_tensor.size(-1), -1)
-            factor = lazy_tensor._matmul(res)
-            factor = factor.view(num_batch, lazy_tensor.size(-2), -1, num_cols)
-            factor = factor.transpose(-3, -2).contiguous().view(-1, num_cols)
-            res = factor.contiguous().view(num_batch, -1, num_cols)
-        else:
-            res = res.view(lazy_tensor.size(-1), -1)
-            factor = lazy_tensor._matmul(res)
-            factor = factor.view(lazy_tensor.size(-2), -1, num_cols).transpose(-3, -2).contiguous().view(-1, num_cols)
-            res = factor.contiguous().view(-1, num_cols)
-
+    for lazy_tensor in lazy_tensors:
+        res = res.view(*lazy_tensor.batch_shape, lazy_tensor.size(-1), -1)
+        factor = lazy_tensor._matmul(res)
+        factor = factor.view(*lazy_tensor.batch_shape, lazy_tensor.size(-2), -1, num_cols).transpose(-3, -2)
+        res = factor.contiguous().view(*lazy_tensor.batch_shape, -1, num_cols)
     return res
 
 
 def _t_matmul(lazy_tensors, rhs):
     res = rhs.contiguous()
-    is_batch = rhs.ndimension() == 3
     num_cols = rhs.size(-1)
-    for lazy_tensor in list(lazy_tensors)[::-1]:
-        if is_batch:
-            num_batch = res.size(0)
-            res = res.view(num_batch, lazy_tensor.size(-2), -1)
-            factor = lazy_tensor._t_matmul(res)
-            factor = factor.view(num_batch, lazy_tensor.size(-1), -1, num_cols)
-            factor = factor.transpose(-3, -2).contiguous().view(-1, num_cols)
-            res = factor.contiguous().view(num_batch, -1, num_cols)
-        else:
-            res = res.view(lazy_tensor.size(-2), -1)
-            factor = lazy_tensor._t_matmul(res)
-            factor = factor.view(lazy_tensor.size(-1), -1, num_cols).transpose(-3, -2).contiguous().view(-1, num_cols)
-            res = factor.contiguous().view(-1, num_cols)
+    for lazy_tensor in lazy_tensors:
+        res = res.view(*lazy_tensor.batch_shape, lazy_tensor.size(-2), -1)
+        factor = lazy_tensor._t_matmul(res)
+        factor = factor.view(*lazy_tensor.batch_shape, lazy_tensor.size(-1), -1, num_cols).transpose(-3, -2)
+        res = factor.contiguous().view(*lazy_tensor.batch_shape, -1, num_cols)
     return res
 
 
@@ -92,103 +73,10 @@ class KroneckerProductLazyTensor(LazyTensor):
             res = res.squeeze(-1)
         return res
 
-    def _quad_form_derivative(self, left_vecs, right_vecs):
-        if left_vecs.ndimension() == 1:
-            left_vecs = left_vecs.unsqueeze(1)
-            right_vecs = right_vecs.unsqueeze(1)
-
-        left_vecs = left_vecs.contiguous()
-        right_vecs = right_vecs.contiguous()
-        is_batch = left_vecs.ndimension() == 3
-        row_sizes = [lazy_tensor.size(-2) for lazy_tensor in self.lazy_tensors]
-        col_sizes = [lazy_tensor.size(-1) for lazy_tensor in self.lazy_tensors]
-
-        res = []
-        if not is_batch:
-            _, s = left_vecs.size()
-
-            # We need to call the derivative routine for each of the sub matrices
-            for i, lazy_tensor in enumerate(self.lazy_tensors):
-                m_left_prev = _prod(row_sizes[i + 1 :])
-                m_left_next = _prod(row_sizes[:i])
-                m_left_i = row_sizes[i]
-                m_right_prev = _prod(col_sizes[i + 1 :])
-                m_right_next = _prod(col_sizes[:i])
-                m_right_i = col_sizes[i]
-
-                # The right vectors are the same for each of the sub matrices
-                # Here we're reshaping them so that the batch dimension corresponds to the correct vectors
-                right_vecs_i = right_vecs.view(m_right_prev, m_right_i, m_right_next, s).transpose(0, 1).contiguous()
-                right_vecs_i = right_vecs_i.view(m_right_i, m_right_prev * m_right_next * s)
-
-                # The left vectors need to be multiplied by all other sub matrices in the kronecker product
-                # First multiply by all the matrices in the KP that are to follow this matrix
-                if i != len(self.lazy_tensors) - 1:
-                    left_vecs_i = left_vecs.view(m_left_prev, m_left_i * m_left_next * s)
-                    left_vecs_i = _t_matmul(self.lazy_tensors[i + 1 :], left_vecs_i)
-                    m_left_prev = m_right_prev
-                    left_vecs_i = left_vecs_i.view(m_left_prev, m_left_i * m_left_next, s)
-                else:
-                    left_vecs_i = left_vecs.view(m_left_prev, m_left_i * m_left_next, s)
-
-                # Next multiply by all the previous matrices in the KP
-                left_vecs_i = left_vecs_i.view(m_left_prev, m_left_i, m_left_next, s).transpose(0, 1)
-                if i != 0:
-                    left_vecs_i = left_vecs_i.transpose(0, 2).contiguous().view(m_left_next, m_left_prev * m_left_i * s)
-                    left_vecs_i = _t_matmul(self.lazy_tensors[:i], left_vecs_i)
-                    m_left_next = m_right_next
-                    left_vecs_i = left_vecs_i.view(m_left_next, m_left_prev, m_left_i, s).transpose(0, 2)
-
-                # Now reshape the left vectors
-                left_vecs_i = left_vecs_i.contiguous().view(m_left_i, m_left_prev * m_left_next * s).contiguous()
-                res = res + list(lazy_tensor._quad_form_derivative(left_vecs_i, right_vecs_i))
-        else:
-            batch_size, _, s = left_vecs.size()
-
-            for i, lazy_tensor in enumerate(self.lazy_tensors):
-                m_right_prev = _prod(col_sizes[i + 1 :])
-                m_right_next = _prod(col_sizes[:i])
-                m_right_i = col_sizes[i]
-                m_left_prev = _prod(row_sizes[i + 1 :])
-                m_left_next = _prod(row_sizes[:i])
-                m_left_i = row_sizes[i]
-
-                right_vecs_i = right_vecs.view(batch_size, m_right_prev, m_right_i, m_right_next, s)
-                right_vecs_i = right_vecs_i.transpose(1, 2).contiguous()
-                right_vecs_i = right_vecs_i.view(batch_size, m_right_i, m_right_prev * m_right_next * s)
-
-                if i != len(self.lazy_tensors) - 1:
-                    left_vecs_i = left_vecs.view(batch_size, m_left_prev, m_left_i * m_left_next * s)
-                    left_vecs_i = _t_matmul(self.lazy_tensors[i + 1 :], left_vecs_i)
-                    m_left_prev = m_right_prev
-                    left_vecs_i = left_vecs_i.contiguous().view(batch_size, m_left_prev, m_left_i * m_left_next, s)
-                else:
-                    left_vecs_i = left_vecs.view(batch_size, m_left_prev, m_left_i * m_left_next, s)
-
-                left_vecs_i = left_vecs_i.view(batch_size, m_left_prev, m_left_i, m_left_next, s)
-                left_vecs_i = left_vecs_i.transpose(1, 2)
-                if i != 0:
-                    left_vecs_i = left_vecs_i.transpose(1, 3).contiguous()
-                    left_vecs_i = left_vecs_i.view(batch_size, m_left_next, m_left_prev * m_left_i * s)
-                    left_vecs_i = _t_matmul(self.lazy_tensors[:i], left_vecs_i)
-                    m_left_next = m_right_next
-                    left_vecs_i = left_vecs_i.view(batch_size, m_left_next, m_left_prev, m_left_i, s)
-                    left_vecs_i = left_vecs_i.transpose(1, 3)
-                left_vecs_i = left_vecs_i.contiguous().view(batch_size, m_left_i, m_left_prev * m_left_next * s)
-                left_vecs_i = left_vecs_i.contiguous()
-                res = res + list(lazy_tensor._quad_form_derivative(left_vecs_i, right_vecs_i))
-
-        return res
-
     def _size(self):
         left_size = _prod(lazy_tensor.size(-2) for lazy_tensor in self.lazy_tensors)
         right_size = _prod(lazy_tensor.size(-1) for lazy_tensor in self.lazy_tensors)
-
-        is_batch = self.lazy_tensors[0].ndimension() == 3
-        if is_batch:
-            return torch.Size((self.lazy_tensors[0].size(0), left_size, right_size))
-        else:
-            return torch.Size((left_size, right_size))
+        return torch.Size((*self.lazy_tensors[0].batch_shape, left_size, right_size))
 
     def _transpose_nonbatch(self):
         return self.__class__(*(lazy_tensor._transpose_nonbatch() for lazy_tensor in self.lazy_tensors), **self._kwargs)
@@ -197,7 +85,7 @@ class KroneckerProductLazyTensor(LazyTensor):
         res = torch.ones(left_indices.size(), dtype=self.dtype, device=self.device)
         left_size = self.size(-2)
         right_size = self.size(-1)
-        for lazy_tensor in list(self.lazy_tensors)[::-1]:
+        for lazy_tensor in self.lazy_tensors:
             left_size = left_size / lazy_tensor.size(-2)
             right_size = right_size / lazy_tensor.size(-1)
             left_indices_i = left_indices.div(left_size)
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index 43ee7f00..5ce14b13 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -212,13 +212,43 @@ class LazyEvaluatedKernelTensor(LazyTensor):
                 self._cached_kernel_eval = self._cached_kernel_eval[0]
             if not isinstance(self._cached_kernel_eval, LazyTensor):
                 self._cached_kernel_eval = NonLazyTensor(self._cached_kernel_eval)
-
             return self._cached_kernel_eval
 
     @cached
     def evaluate(self):
         return self.evaluate_kernel().evaluate()
 
+    def exact_predictive_mean(
+        self,
+        full_mean,
+        train_inputs,
+        train_labels,
+        num_train,
+        likelihood,
+        precomputed_cache=None,
+        non_batch_train=False,
+    ):
+        if self.kernel.has_custom_exact_predictions:
+            return self.evaluate_kernel().exact_predictive_mean(
+                full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache, non_batch_train
+            )
+        else:
+            return super(LazyEvaluatedKernelTensor, self).exact_predictive_mean(
+                full_mean, train_inputs, train_labels, num_train, likelihood, precomputed_cache, non_batch_train
+            )
+
+    def exact_predictive_covar(
+        self, train_inputs, num_train, likelihood, precomputed_cache=None, non_batch_train=False
+    ):
+        if self.kernel.has_custom_exact_predictions:
+            return self.evaluate_kernel().exact_predictive_covar(
+                train_inputs, num_train, likelihood, precomputed_cache, non_batch_train
+            )
+        else:
+            return super(LazyEvaluatedKernelTensor, self).exact_predictive_covar(
+                train_inputs, num_train, likelihood, precomputed_cache, non_batch_train
+            )
+
     def repeat(self, *sizes):
         if self.squeeze_row or self.squeeze_col:
             raise RuntimeError("Can't repeat a row/col of a LazyEvaluatedKernelTensor")
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index f4514df5..3740cb8d 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -5,7 +5,7 @@ import math
 import gpytorch
 import torch
 
-from .. import settings
+from .. import beta_features, settings
 from ..functions._inv_matmul import InvMatmul
 from ..functions._inv_quad_log_det import InvQuadLogDet
 from ..functions._matmul import Matmul
@@ -144,6 +144,36 @@ class LazyTensor(object):
         """
         return self.diag()
 
+    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
+        """
+        Computes a cache for K_X*X (K_XX + sigma^2 I)^-1 K_X*X if possible. By default, this does no work and returns
+        the first argument.
+
+        Args:
+            train_train_covar_inv_root (:obj:`torch.tensor`): a root of (K_XX + sigma^2 I)^-1
+            test_train_covar (:obj:`torch.tensor`): the observed noise (from the likelihood)
+
+        Returns
+            - A precomputed cache
+        """
+        return train_train_covar_inv_root.detach()
+
+    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
+        """
+        Computes :math:`K_{X^{*}X} S` given a precomputed cache
+        Where :math:`S` is a tensor such that :math:`SS^{\\top} = (K_{XX} + \sigma^2 I)^{-1}`
+
+        Args:
+            precomputed_cache (:obj:`torch.tensor`): What was computed in _exact_predictive_covar_inv_quad_form_cache
+            test_train_covar (:obj:`torch.tensor`): The observed noise (from the likelihood)
+
+        Returns
+            :obj:`~gpytorch.lazy.LazyTensor`: :math:`K_{X^{*}X} S`
+        """
+        # Here the precomputed cache represents S,
+        # where S S^T = (K_XX + sigma^2 I)^-1
+        return test_train_covar.matmul(precomputed_cache)
+
     def _getitem(self, *indices):
         """
         Supports subindexing of the matrix this LazyTensor represents. This may return either another
@@ -359,22 +389,28 @@ class LazyTensor(object):
             :obj:`torch.tensor`: derivative with respect to the arguments that are actually used to represent this
                                    this LazyTensor.
         """
-        args = self.representation()
+        from collections import deque
 
-        toggled = [False] * len(args)
+        args = tuple(self.representation())
+        args_with_grads = tuple(arg for arg in args if arg.requires_grad)
 
-        for i, arg in enumerate(args):
-            if not arg.requires_grad:
-                arg.requires_grad = True
-                toggled[i] = True
+        # Easy case: if we don't require any gradients, then just return!
+        if not len(args_with_grads):
+            return tuple(None for _ in args)
 
-        loss = (left_vecs * self._matmul(right_vecs)).sum()
-        loss.requires_grad_(True)
-        grads = torch.autograd.grad(loss, args, allow_unused=True)
+        # Normal case: we'll use the autograd to get us a derivative
+        with torch.autograd.enable_grad():
+            loss = (left_vecs * self._matmul(right_vecs)).sum()
+            loss.requires_grad_(True)
+            actual_grads = deque(torch.autograd.grad(loss, args_with_grads, allow_unused=True))
 
-        for i, arg in enumerate(args):
-            if toggled[i]:
-                arg.requires_grad = False
+        # Now make sure that the object we return has one entry for every item in args
+        grads = []
+        for arg in args:
+            if arg.requires_grad:
+                grads.append(actual_grads.popleft())
+            else:
+                grads.append(None)
 
         return grads
 
@@ -611,6 +647,124 @@ class LazyTensor(object):
         """
         return self.representation_tree()(*self.representation())
 
+    def exact_predictive_mean(
+        self,
+        full_mean,
+        train_inputs,
+        train_labels,
+        num_train,
+        likelihood,
+        precomputed_cache=None,
+        non_batch_train=False,
+    ):
+        """
+        Computes the posterior predictive covariance of a GP
+        Assumes that self is the block prior covariance matrix of training and testing points
+        [ K_XX, K_XX*; K_X*X, K_X*X* ]
+
+        Args:
+            full_mean (:obj:`torch.tensor`): the training and test prior means, stacked on top of each other
+            train_inputs (:obj:`torch.tensor`): The training data inputs
+            train_labels (:obj:`torch.tensor`): the training labels minus the training prior mean
+            noise (:obj:`torch.tensor`): the observed noise (from the likelihood)
+            precomputed_cache (optional): speeds up subsequent computations (default: None)
+
+        Returns:
+            :obj:`torch.tensor`: The predictive posterior mean of the test points
+        """
+        from ..distributions import MultivariateNormal
+
+        if precomputed_cache is None:
+            train_mean = full_mean.narrow(-1, 0, num_train)
+            if non_batch_train and self.dim() == 3:
+                train_train_covar = self[0, :num_train, :num_train]
+            else:
+                train_train_covar = self[..., :num_train, :num_train]
+
+            train_mean = full_mean.narrow(-1, 0, train_train_covar.size(-1))
+            if non_batch_train and train_mean.dim() == 2:
+                train_mean = train_mean[0]
+                train_labels = train_labels[0]
+            mvn = likelihood(MultivariateNormal(train_mean, train_train_covar), train_inputs)
+
+            train_mean, train_train_covar = mvn.mean, mvn.lazy_covariance_matrix
+
+            train_labels_offset = train_labels - train_mean
+
+            if self.dim() == 3:
+                # Batch mode
+                train_labels_offset = train_labels_offset.unsqueeze(-1)
+                precomputed_cache = train_train_covar.inv_matmul(train_labels_offset).squeeze(-1)
+            else:
+                # Standard mode
+                precomputed_cache = train_train_covar.inv_matmul(train_labels_offset)
+
+        test_mean = full_mean.narrow(-1, train_labels.size(-1), full_mean.size(-1) - train_labels.size(-1))
+
+        if self.dim() == 3:
+            test_train_covar = self[:, num_train:, :num_train]
+            res = test_train_covar.matmul(precomputed_cache.unsqueeze(-1)).squeeze(-1)
+        else:
+            test_train_covar = self[num_train:, :num_train]
+            if non_batch_train and precomputed_cache.dim() == 2:
+                precomputed_cache = precomputed_cache[0]
+            res = test_train_covar.matmul(precomputed_cache)
+
+        res = res + test_mean
+
+        return res, precomputed_cache.detach()
+
+    def exact_predictive_covar(
+        self, train_inputs, num_train, likelihood, precomputed_cache=None, non_batch_train=False
+    ):
+        """
+        Computes the posterior predictive covariance of a GP
+        Assumes that self is the block prior covariance matrix of training and testing points
+        [ K_XX, K_XX*; K_X*X, K_X*X* ]
+
+        Args:
+            train_inputs (:obj:`torch.tensor`): The training data inputs
+            num_train (int): The number of training points in the full covariance matrix
+            noise (scalar): The observed noise (from the likelihood)
+            precomputed_cache (optional): speeds up subsequent computations (default: None)
+
+        Returns:
+            :obj:`gpytorch.lazy.LazyTensor`: A LazyTensor representing the predictive posterior covariance of the
+                                               test points
+        """
+        from ..distributions import MultivariateNormal
+
+        train_train_covar = self[..., :num_train, :num_train]
+        test_train_covar = self[..., num_train:, :num_train]
+        test_test_covar = self[..., num_train:, num_train:]
+
+        train_train_covar = likelihood(
+            MultivariateNormal(torch.zeros(1), train_train_covar), train_inputs
+        ).lazy_covariance_matrix
+        if not beta_features.fast_pred_var.on():
+            from .matmul_lazy_tensor import MatmulLazyTensor
+
+            test_train_covar = test_train_covar.evaluate()
+            train_test_covar = test_train_covar.transpose(-1, -2)
+            covar_correction_rhs = train_train_covar.inv_matmul(train_test_covar).mul(-1)
+            res = test_test_covar + MatmulLazyTensor(test_train_covar, covar_correction_rhs)
+            return res, None
+
+        if precomputed_cache is None:
+            if non_batch_train and train_train_covar.dim() == 3:
+                train_train_covar_inv_root = train_train_covar[0].root_inv_decomposition().root.evaluate()
+            else:
+                train_train_covar_inv_root = train_train_covar.root_inv_decomposition().root.evaluate()
+            precomputed_cache = self._exact_predictive_covar_inv_quad_form_cache(
+                train_train_covar_inv_root, test_train_covar
+            )
+
+        from .root_lazy_tensor import RootLazyTensor
+
+        covar_inv_quad_form_root = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache, test_train_covar)
+        res = test_test_covar + RootLazyTensor(covar_inv_quad_form_root).mul(-1)
+        return res, precomputed_cache
+
     def inv_matmul(self, tensor):
         """
         Computes a linear solve (w.r.t self = :math:`K`) with several right hand sides :math:`M`.
diff --git a/gpytorch/lazy/sum_lazy_tensor.py b/gpytorch/lazy/sum_lazy_tensor.py
index 9e3a0538..6d9b3777 100644
--- a/gpytorch/lazy/sum_lazy_tensor.py
+++ b/gpytorch/lazy/sum_lazy_tensor.py
@@ -51,6 +51,24 @@ class SumLazyTensor(LazyTensor):
         lazy_tensors_t = [lazy_tensor.transpose(-1, -2) for lazy_tensor in self.lazy_tensors]
         return SumLazyTensor(*lazy_tensors_t)
 
+    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
+        return tuple(
+            lazy_tensor._exact_predictive_covar_inv_quad_form_cache(
+                train_train_covar_inv_root, test_train_covar_comp
+            ).detach()
+            for lazy_tensor, test_train_covar_comp in zip(self.lazy_tensors, test_train_covar.lazy_tensors)
+        )
+
+    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
+        # Here the precomputed cache is a list
+        # where each component in the list is the precomputed cache for each component lazy tensor
+        return sum(
+            lazy_tensor._exact_predictive_covar_inv_quad_form_root(cache_comp, test_train_covar_comp)
+            for lazy_tensor, cache_comp, test_train_covar_comp in zip(
+                self.lazy_tensors, precomputed_cache, test_train_covar.lazy_tensors
+            )
+        )
+
     @cached
     def evaluate(self):
         return sum(lazy_tensor.evaluate() for lazy_tensor in self.lazy_tensors)
diff --git a/gpytorch/likelihoods/multitask_gaussian_likelihood.py b/gpytorch/likelihoods/multitask_gaussian_likelihood.py
index b8168053..8c58a551 100644
--- a/gpytorch/likelihoods/multitask_gaussian_likelihood.py
+++ b/gpytorch/likelihoods/multitask_gaussian_likelihood.py
@@ -323,7 +323,7 @@ class MultitaskGaussianLikelihoodKronecker(_MultitaskGaussianLikelihoodBase):
             if task_var_lt.size(0) == 1:
                 task_var_lt = task_var_lt.repeat(eye_lt.size(0), 1, 1)
 
-        covar_kron_lt = KroneckerProductLazyTensor(task_var_lt, eye_lt)
+        covar_kron_lt = KroneckerProductLazyTensor(eye_lt, task_var_lt)
         covar = covar + covar_kron_lt
 
         noise = self.noise
diff --git a/gpytorch/likelihoods/softmax_likelihood.py b/gpytorch/likelihoods/softmax_likelihood.py
index 3137d7ef..55b5d264 100755
--- a/gpytorch/likelihoods/softmax_likelihood.py
+++ b/gpytorch/likelihoods/softmax_likelihood.py
@@ -32,8 +32,6 @@ class SoftmaxLikelihood(Likelihood):
 
         n_samples = settings.num_likelihood_samples.value()
         samples = latent_func.rsample(sample_shape=torch.Size((n_samples,)))
-        if samples.dim() == 2:
-            samples = samples.unsqueeze(-1).transpose(-2, -1)
         samples = samples.permute(1, 2, 0).contiguous()  # Now n_featuers, n_data, n_samples
         if samples.ndimension() != 3:
             raise RuntimeError("f should have 3 dimensions: features x data x samples")
diff --git a/gpytorch/models/exact_gp.py b/gpytorch/models/exact_gp.py
index bae82207..c9172646 100644
--- a/gpytorch/models/exact_gp.py
+++ b/gpytorch/models/exact_gp.py
@@ -2,11 +2,11 @@
 
 import warnings
 import torch
+from ..functions import exact_predictive_mean, exact_predictive_covar
 from ..distributions import MultivariateNormal, MultitaskMultivariateNormal
 from ..likelihoods import _GaussianLikelihoodBase
 from .. import settings
 from .gp import GP
-from .exact_prediction_strategies import prediction_strategy
 
 
 class ExactGP(GP):
@@ -115,7 +115,6 @@ class ExactGP(GP):
 
             num_train = 0
             train_targets = None
-
             if self.train_targets is not None:
                 train_targets = self.train_targets
 
@@ -148,44 +147,33 @@ class ExactGP(GP):
                     # non-batch mode (standard)
                     num_train = train_targets.size(-1)
 
-                train_train_covar = full_covar[..., :num_train, :num_train].evaluate_kernel()
-                test_train_covar = full_covar[..., num_train:, :num_train]
-                test_test_covar = full_covar[..., num_train:, num_train:]
-
-                pred_strat = prediction_strategy(train_train_covar)
-
-                predictive_mean, mean_cache = pred_strat.exact_predictive_mean(
-                    test_train_covar=test_train_covar,
-                    full_mean=full_mean,
-                    train_inputs=train_inputs,
-                    train_labels=train_targets,
-                    num_train=num_train,
-                    likelihood=self.likelihood,
-                    precomputed_cache=self.mean_cache,
-                    non_batch_train=non_batch_train,
-                )
-
-                predictive_covar, covar_cache = pred_strat.exact_predictive_covar(
-                    test_train_covar=test_train_covar,
-                    test_test_covar=test_test_covar,
-                    train_inputs=train_inputs,
-                    num_train=num_train,
-                    likelihood=self.likelihood,
-                    precomputed_cache=self.covar_cache,
-                    non_batch_train=non_batch_train,
-                )
-
-                self.mean_cache = mean_cache
-                self.covar_cache = covar_cache
-                if num_tasks > 1:
-                    if train_targets.ndimension() == 2:
-                        # Batch multitask
-                        predictive_mean = predictive_mean.view(train_targets.size(0), -1, num_tasks).contiguous()
-                    else:
-                        # Standard multitask
-                        predictive_mean = predictive_mean.view(-1, num_tasks).contiguous()
-            else:
-                predictive_mean = full_mean
-                predictive_covar = full_covar
+            predictive_mean, mean_cache = exact_predictive_mean(
+                full_covar=full_covar,
+                full_mean=full_mean,
+                train_inputs=train_inputs,
+                train_labels=train_targets,
+                num_train=num_train,
+                likelihood=self.likelihood,
+                precomputed_cache=self.mean_cache,
+                non_batch_train=non_batch_train,
+            )
+            predictive_covar, covar_cache = exact_predictive_covar(
+                full_covar=full_covar,
+                train_inputs=train_inputs,
+                num_train=num_train,
+                likelihood=self.likelihood,
+                precomputed_cache=self.covar_cache,
+                non_batch_train=non_batch_train,
+            )
+
+            self.mean_cache = mean_cache
+            self.covar_cache = covar_cache
+            if num_tasks > 1:
+                if train_targets.ndimension() == 2:
+                    # Batch multitask
+                    predictive_mean = predictive_mean.view(train_targets.size(0), -1, num_tasks).contiguous()
+                else:
+                    # Standard multitask
+                    predictive_mean = predictive_mean.view(-1, num_tasks).contiguous()
 
             return full_output.__class__(predictive_mean, predictive_covar)
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
deleted file mode 100644
index 84f6f467..00000000
--- a/gpytorch/models/exact_prediction_strategies.py
+++ /dev/null
@@ -1,393 +0,0 @@
-import torch
-from .. import beta_features
-from ..lazy import LazyTensor, InterpolatedLazyTensor, SumLazyTensor, MatmulLazyTensor, RootLazyTensor
-from ..utils.interpolation import left_interp, left_t_interp
-
-_PREDICTION_STRATEGY_REGISTRY = {}
-
-
-def register_prediction_strategy(lazy_tsr_type):
-    if not isinstance(lazy_tsr_type, type) and issubclass(lazy_tsr_type, LazyTensor):
-        raise TypeError(f"register_prediction_strategy expects a LazyTensor subtype but got {lazy_tsr_type}")
-
-    def decorator(cls):
-        _PREDICTION_STRATEGY_REGISTRY[lazy_tsr_type] = cls
-        return cls
-
-    return decorator
-
-
-def prediction_strategy(train_train_covar):
-    cls = _PREDICTION_STRATEGY_REGISTRY.get(type(train_train_covar), DefaultPredictionStrategy)
-    return cls(train_train_covar)
-
-
-class DefaultPredictionStrategy(object):
-    def __init__(self, train_train_covar):
-        self.train_train_covar = train_train_covar
-
-    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
-        """
-        Computes a cache for K_X*X (K_XX + sigma^2 I)^-1 K_X*X if possible. By default, this does no work and returns
-        the first argument.
-
-        Args:
-            train_train_covar_inv_root (:obj:`torch.tensor`): a root of (K_XX + sigma^2 I)^-1
-            test_train_covar (:obj:`torch.tensor`): the observed noise (from the likelihood)
-
-        Returns
-            - A precomputed cache
-        """
-        return train_train_covar_inv_root.detach()
-
-    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
-        """
-        Computes :math:`K_{X^{*}X} S` given a precomputed cache
-        Where :math:`S` is a tensor such that :math:`SS^{\\top} = (K_{XX} + \sigma^2 I)^{-1}`
-
-        Args:
-            precomputed_cache (:obj:`torch.tensor`): What was computed in _exact_predictive_covar_inv_quad_form_cache
-            test_train_covar (:obj:`torch.tensor`): The observed noise (from the likelihood)
-
-        Returns
-            :obj:`~gpytorch.lazy.LazyTensor`: :math:`K_{X^{*}X} S`
-        """
-        # Here the precomputed cache represents S,
-        # where S S^T = (K_XX + sigma^2 I)^-1
-        return test_train_covar.matmul(precomputed_cache)
-
-    def exact_predictive_mean(
-        self,
-        test_train_covar,
-        full_mean,
-        train_inputs,
-        train_labels,
-        num_train,
-        likelihood,
-        precomputed_cache=None,
-        non_batch_train=False,
-    ):
-        """
-        Computes the posterior predictive covariance of a GP
-
-        Args:
-            test_train_covar (:obj:`gpytorch.lazy.LazyTensor`): Covariance matrix between test and train inputs
-            full_mean (:obj:`torch.tensor`): the training and test prior means, stacked on top of each other
-            train_inputs (:obj:`torch.tensor`): The training data inputs
-            train_labels (:obj:`torch.tensor`): the training labels minus the training prior mean
-            noise (:obj:`torch.tensor`): the observed noise (from the likelihood)
-            precomputed_cache (optional): speeds up subsequent computations (default: None)
-
-        Returns:
-            :obj:`torch.tensor`: The predictive posterior mean of the test points
-        """
-        from ..distributions import MultivariateNormal
-
-        if precomputed_cache is None:
-            train_mean = full_mean.narrow(-1, 0, num_train)
-
-            if non_batch_train and self.train_train_covar.dim() == 3:
-                train_train_covar = self.train_train_covar[0]
-            else:
-                train_train_covar = self.train_train_covar
-
-            train_mean = full_mean.narrow(-1, 0, train_train_covar.size(-1))
-            if non_batch_train and train_mean.dim() == 2:
-                train_mean = train_mean[0]
-                train_labels = train_labels[0]
-            mvn = likelihood(MultivariateNormal(train_mean, train_train_covar), train_inputs)
-
-            train_mean, train_train_covar = mvn.mean, mvn.lazy_covariance_matrix
-
-            train_labels_offset = train_labels - train_mean
-
-            if self.train_train_covar.dim() == 3:
-                # Batch mode
-                train_labels_offset = train_labels_offset.unsqueeze(-1)
-                precomputed_cache = train_train_covar.inv_matmul(train_labels_offset).squeeze(-1)
-            else:
-                # Standard mode
-                precomputed_cache = train_train_covar.inv_matmul(train_labels_offset)
-
-        test_mean = full_mean.narrow(-1, train_labels.size(-1), full_mean.size(-1) - train_labels.size(-1))
-
-        if self.train_train_covar.dim() == 3:
-            res = test_train_covar.matmul(precomputed_cache.unsqueeze(-1)).squeeze(-1)
-        else:
-            if non_batch_train and precomputed_cache.dim() == 2:
-                precomputed_cache = precomputed_cache[0]
-            res = test_train_covar.matmul(precomputed_cache)
-
-        res = res + test_mean
-
-        return res, precomputed_cache.detach()
-
-    def exact_predictive_covar(
-        self,
-        test_train_covar,
-        test_test_covar,
-        train_inputs,
-        num_train,
-        likelihood,
-        precomputed_cache=None,
-        non_batch_train=False,
-    ):
-        """
-        Computes the posterior predictive covariance of a GP
-
-        Args:
-            test_train_covar (:obj:`gpytorch.lazy.LazyTensor`): Covariance matrix between test and train inputs
-            test_test_covar (:obj:`gpytorch.lazy.LazyTensor`): Covariance matrix between test inputs
-            train_inputs (:obj:`torch.tensor`): The training data inputs
-            num_train (int): The number of training points in the full covariance matrix
-            noise (scalar): The observed noise (from the likelihood)
-            precomputed_cache (optional): speeds up subsequent computations (default: None)
-            non_batch_train (bool, optional): True if the training data was not batch mode.
-
-        Returns:
-            :obj:`gpytorch.lazy.LazyTensor`: A LazyTensor representing the predictive posterior covariance of the
-                                               test points
-        """
-        from ..distributions import MultivariateNormal
-
-        train_train_covar = likelihood(
-            MultivariateNormal(torch.zeros(1), self.train_train_covar), train_inputs
-        ).lazy_covariance_matrix
-        if not beta_features.fast_pred_var.on():
-            test_train_covar = test_train_covar.evaluate()
-            train_test_covar = test_train_covar.transpose(-1, -2)
-            covar_correction_rhs = train_train_covar.inv_matmul(train_test_covar).mul(-1)
-            res = test_test_covar + MatmulLazyTensor(test_train_covar, covar_correction_rhs)
-            return res, None
-
-        if precomputed_cache is None:
-            if non_batch_train and train_train_covar.dim() == 3:
-                train_train_covar_inv_root = train_train_covar[0].root_inv_decomposition().root.evaluate()
-            else:
-                train_train_covar_inv_root = train_train_covar.root_inv_decomposition().root.evaluate()
-            precomputed_cache = self._exact_predictive_covar_inv_quad_form_cache(
-                train_train_covar_inv_root, test_train_covar
-            )
-
-        covar_inv_quad_form_root = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache, test_train_covar)
-        res = test_test_covar + RootLazyTensor(covar_inv_quad_form_root).mul(-1)
-        return res, precomputed_cache
-
-
-@register_prediction_strategy(InterpolatedLazyTensor)
-class InterpolatedPredictionStrategy(DefaultPredictionStrategy):
-    def __init__(self, train_train_covar):
-        if not isinstance(train_train_covar, InterpolatedLazyTensor):
-            raise TypeError(
-                "InterpolatedPredictionStrategy can only be used with InterpolatedLazyTensors, "
-                f"but got {type(train_train_covar)}"
-            )
-
-        super(InterpolatedPredictionStrategy, self).__init__(train_train_covar)
-
-    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
-        train_interp_indices = test_train_covar.right_interp_indices
-        train_interp_values = test_train_covar.right_interp_values
-        base_lazy_tensor = test_train_covar.base_lazy_tensor
-        base_size = base_lazy_tensor.size(-1)
-        res = base_lazy_tensor.matmul(
-            left_t_interp(train_interp_indices, train_interp_values, train_train_covar_inv_root, base_size)
-        )
-        return res
-
-    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
-        # Here the precomputed cache represents K_UU W S,
-        # where S S^T = (K_XX + sigma^2 I)^-1
-        test_interp_indices = test_train_covar.left_interp_indices
-        test_interp_values = test_train_covar.left_interp_values
-        res = left_interp(test_interp_indices, test_interp_values, precomputed_cache)
-        return res
-
-    def exact_predictive_mean(
-        self,
-        test_train_covar,
-        full_mean,
-        train_inputs,
-        train_labels,
-        num_train,
-        likelihood,
-        precomputed_cache=None,
-        non_batch_train=False,
-    ):
-        from ..distributions import MultivariateNormal
-
-        if precomputed_cache is None:
-            train_interp_indices = self.train_train_covar.left_interp_indices
-            train_interp_values = self.train_train_covar.left_interp_values
-
-            train_mean = full_mean.narrow(-1, 0, train_labels.size(-1))
-
-            mvn = likelihood(MultivariateNormal(train_mean, self.train_train_covar), train_inputs)
-            train_mean, train_train_covar = mvn.mean, mvn.lazy_covariance_matrix
-
-            train_train_covar_inv_labels = train_train_covar.inv_matmul((train_labels - train_mean).unsqueeze(-1))
-
-            # New root factor
-            base_size = self.train_train_covar.base_lazy_tensor.size(-1)
-            precomputed_cache = self.train_train_covar.base_lazy_tensor.matmul(
-                left_t_interp(train_interp_indices, train_interp_values, train_train_covar_inv_labels, base_size)
-            )
-
-            # Prevent backprop through this variable
-            precomputed_cache = precomputed_cache.detach()
-
-        # Compute the exact predictive posterior
-        n_test = test_train_covar.size(-2)
-
-        test_train_covar = test_train_covar.evaluate_kernel()
-
-        test_mean = full_mean.narrow(-1, num_train, n_test)
-        test_interp_indices = test_train_covar.left_interp_indices
-        test_interp_values = test_train_covar.left_interp_values
-        res = left_interp(test_interp_indices, test_interp_values, precomputed_cache).squeeze(-1) + test_mean
-        return res, precomputed_cache
-
-    def exact_predictive_covar(
-        self,
-        test_train_covar,
-        test_test_covar,
-        train_inputs,
-        num_train,
-        likelihood,
-        precomputed_cache=None,
-        non_batch_train=False,
-    ):
-        from ..distributions import MultivariateNormal
-
-        if not beta_features.fast_pred_var.on() and not beta_features.fast_pred_samples.on():
-            return super(InterpolatedPredictionStrategy, self).exact_predictive_covar(
-                test_train_covar, test_test_covar, train_inputs, num_train, likelihood, precomputed_cache
-            )
-
-        train_interp_indices = self.train_train_covar.left_interp_indices
-        train_interp_values = self.train_train_covar.left_interp_values
-
-        test_train_covar = test_train_covar.evaluate_kernel()
-        test_test_covar = test_test_covar.evaluate_kernel()
-
-        test_interp_indices = test_train_covar.left_interp_indices
-        test_interp_values = test_train_covar.left_interp_values
-
-        if (
-            precomputed_cache is None
-            or (beta_features.fast_pred_samples.on() and precomputed_cache[0] is None)
-            or (not beta_features.fast_pred_samples.on() and precomputed_cache[1] is None)
-        ):
-            # Get inverse root
-
-            grv = MultivariateNormal(torch.zeros(1), self.train_train_covar)
-            train_train_covar = likelihood(grv, train_inputs).lazy_covariance_matrix
-
-            # Get probe vectors for inverse root
-            num_probe_vectors = beta_features.fast_pred_var.num_probe_vectors()
-            batch_size = train_interp_indices.size(0)
-            n_inducing = self.train_train_covar.base_lazy_tensor.size(-1)
-            vector_indices = torch.randperm(n_inducing).type_as(train_interp_indices)
-            probe_vector_indices = vector_indices[:num_probe_vectors]
-            test_vector_indices = vector_indices[num_probe_vectors : 2 * num_probe_vectors]
-
-            probe_interp_indices = probe_vector_indices.unsqueeze(1)
-            probe_test_interp_indices = test_vector_indices.unsqueeze(1)
-            dtype = self.train_train_covar.dtype
-            device = self.train_train_covar.device
-            probe_interp_values = torch.ones(num_probe_vectors, 1, dtype=dtype, device=device)
-            if train_interp_indices.ndimension() == 3:
-                probe_interp_indices = probe_interp_indices.unsqueeze(0).expand(batch_size, num_probe_vectors, 1)
-                probe_test_interp_indices = probe_test_interp_indices.unsqueeze(0)
-                probe_test_interp_indices = probe_test_interp_indices.expand(batch_size, num_probe_vectors, 1)
-                probe_interp_values = probe_interp_values.unsqueeze(0).expand(batch_size, num_probe_vectors, 1)
-
-            probe_vectors = InterpolatedLazyTensor(
-                self.train_train_covar.base_lazy_tensor,
-                train_interp_indices,
-                train_interp_values,
-                probe_interp_indices,
-                probe_interp_values,
-            ).evaluate()
-            test_vectors = InterpolatedLazyTensor(
-                self.train_train_covar.base_lazy_tensor,
-                train_interp_indices,
-                train_interp_values,
-                probe_test_interp_indices,
-                probe_interp_values,
-            ).evaluate()
-
-            # Get inverse root
-            train_train_covar_inv_root = train_train_covar.root_inv_decomposition(probe_vectors, test_vectors).root
-            train_train_covar_inv_root = train_train_covar_inv_root.evaluate()
-
-            # New root factor
-            root = self._exact_predictive_covar_inv_quad_form_cache(train_train_covar_inv_root, test_train_covar)
-
-            # Precomputed factor
-            if beta_features.fast_pred_samples.on():
-                inside = self.train_train_covar.base_lazy_tensor + RootLazyTensor(root).mul(-1)
-                inside_root = inside.root_decomposition().root.evaluate()
-                # Prevent backprop through this variable
-                inside_root = inside_root.detach()
-                precomputed_cache = inside_root, None
-            else:
-                # Prevent backprop through this variable
-                root = root.detach()
-                precomputed_cache = None, root
-
-        # Compute the exact predictive posterior
-        if beta_features.fast_pred_samples.on():
-            res = self._exact_predictive_covar_inv_quad_form_root(precomputed_cache[0], test_train_covar)
-            res = RootLazyTensor(res)
-        else:
-            root = left_interp(test_interp_indices, test_interp_values, precomputed_cache[1])
-            res = test_test_covar + RootLazyTensor(root).mul(-1)
-        return res, precomputed_cache
-
-
-@register_prediction_strategy(SumLazyTensor)
-class SumPredictionStrategy(DefaultPredictionStrategy):
-    def __init__(self, train_train_covar):
-        if not isinstance(train_train_covar, SumLazyTensor):
-            raise TypeError(
-                f"SumPredictionStrategy can only be used with SumLazyTensors, but got {type(train_train_covar)}"
-            )
-
-        super(SumPredictionStrategy, self).__init__(train_train_covar)
-
-    def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root, test_train_covar):
-        test_train_covar = test_train_covar.evaluate_kernel()
-        if not isinstance(test_train_covar, SumLazyTensor):
-            return super(SumPredictionStrategy, self)._exact_predictive_covar_inv_quad_form_cache(
-                train_train_covar_inv_root, test_train_covar
-            )
-        else:
-            return tuple(
-                prediction_strategy(lazy_tensor)._exact_predictive_covar_inv_quad_form_cache(
-                    train_train_covar_inv_root, test_train_covar_comp
-                )
-                for lazy_tensor, test_train_covar_comp in zip(
-                    self.train_train_covar.lazy_tensors, test_train_covar.lazy_tensors
-                )
-            )
-
-    def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
-        # Here the precomputed cache is a list
-        # where each component in the list is the precomputed cache for each component lazy tensor
-        test_train_covar = test_train_covar.evaluate_kernel()
-        if not isinstance(test_train_covar, SumLazyTensor):
-            return super(SumPredictionStrategy, self)._exact_predictive_covar_inv_quad_form_root(
-                precomputed_cache, test_train_covar
-            )
-        else:
-            return sum(
-                prediction_strategy(lazy_tensor)._exact_predictive_covar_inv_quad_form_root(
-                    cache_comp, test_train_covar_comp
-                )
-                for lazy_tensor, cache_comp, test_train_covar_comp in zip(
-                    self.train_train_covar.lazy_tensors,
-                    precomputed_cache,
-                    test_train_covar.evaluate_kernel().lazy_tensors,
-                )
-            )


Test patch:
diff --git a/test/lazy/test_kronecker_product_lazy_tensor.py b/test/lazy/test_kronecker_product_lazy_tensor.py
index 4faa113c..8ef40354 100644
--- a/test/lazy/test_kronecker_product_lazy_tensor.py
+++ b/test/lazy/test_kronecker_product_lazy_tensor.py
@@ -8,20 +8,12 @@ from test.lazy._lazy_tensor_test_case import LazyTensorTestCase, RectangularLazy
 
 def kron(a, b):
     res = []
-    if b.ndimension() == 2:
-        for i in range(b.size(0)):
-            row_res = []
-            for j in range(b.size(1)):
-                row_res.append(a * b[i, j])
-            res.append(torch.cat(row_res, 1))
-        return torch.cat(res, 0)
-    else:
-        for i in range(b.size(1)):
-            row_res = []
-            for j in range(b.size(2)):
-                row_res.append(a * b[:, i, j].unsqueeze(1).unsqueeze(2))
-            res.append(torch.cat(row_res, 2))
-        return torch.cat(res, 1)
+    for i in range(a.size(-2)):
+        row_res = []
+        for j in range(a.size(-1)):
+            row_res.append(b * a[..., i, j].unsqueeze(-1).unsqueeze(-2))
+        res.append(torch.cat(row_res, -1))
+    return torch.cat(res, -2)
 
 
 class TestKroneckerProductLazyTensor(LazyTensorTestCase, unittest.TestCase):