Repository: cornellius-gp/gpytorch
Pull request #307 (2018-10-01 16:47:11)

Kernels now take `diag, dim_groups` as arguments (fixes issues with ARD kernels)
This should fix several issues we have with ARD batch kernels.

TLDR:
- `Kernel.forward` now takes 4 arguments: `x1`, `x2`, `diag`, and `dim_groups`
  - `diag` tells the kernel to be computed in diagonal mode (variances only, no covariances)
  - `dim_groups=<k>` computes k batches of kernels: each operating on one of `k` subsets of dimensions. (This is to be used with additive kernels, SKI, etc.)
- For most kernel implementations - you won't have to worry about these arguments much. They just need to be passed into `_create_input_grid`

The motivation for this: for kernel diagonals and additive kernels, we have been relying on hacking the batch dimension. For example, computing the kernel diagonal of a `n x d` input can be done easily if we make the input `n x 1 x d`. Similarly, additive kernels can be efficiently computed if we make the `n x d` input a `d x n x 1` input and then sum along the batch dimension. This works well if we don't have ARD or batch-specific parameters, but now we do and it causes breaking issues. (See #249)

To make sure that this isn't an issue, but also keep everything easily programmable:
- We won't do batch-dimension hacking anymore
- Computing kernel diagonals/additive kernels is now done explicitly by passing in the `dim`/`dim_groups` arguments
- For most kernels, this won't require any extra code. Most kernels now use the `_create_input_grid` helper, which automatically will do the correct thing for the diag/dim_groups modes.

I've included some documentation in the Kernel base_class to describe this for people who want to code up their own kernels.

This is currently a WIP - hopefully will be done in ~4 hours.

Balandat: So I'm all for being explicit and not doing a bunch of magic, my only worry is that by adding additional args we are making it harder / less accessible for people to write their own kernels. From reading the docstring, I'm not really clear what `dim_groups` actually does... Is there a way to simplify this slightly? Not sure if defining this as `forward(slef, x1, x2, **kwargs)` is any better in that regard...
jacobrgardner: @Balandat Geoff and I have been talking about this exact concern in the past couple of hours. I generally had the same concern. Geoff pointed out that in most cases those args just get passed directly on to a helper method. I do agree we can make it clearer exactly what the args do in the docstrings though 

Patch:
diff --git a/gpytorch/kernels/__init__.py b/gpytorch/kernels/__init__.py
index cff952b8..906c9ffc 100644
--- a/gpytorch/kernels/__init__.py
+++ b/gpytorch/kernels/__init__.py
@@ -4,7 +4,7 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 from .kernel import Kernel, AdditiveKernel, ProductKernel
-from .additive_grid_interpolation_kernel import AdditiveGridInterpolationKernel
+from .additive_structure_kernel import AdditiveStructureKernel
 from .cosine_kernel import CosineKernel
 from .grid_kernel import GridKernel
 from .grid_interpolation_kernel import GridInterpolationKernel
@@ -13,9 +13,9 @@ from .inducing_point_kernel import InducingPointKernel
 from .lcm_kernel import LCMKernel
 from .linear_kernel import LinearKernel
 from .matern_kernel import MaternKernel
-from .multiplicative_grid_interpolation_kernel import MultiplicativeGridInterpolationKernel
 from .multitask_kernel import MultitaskKernel
 from .periodic_kernel import PeriodicKernel
+from .product_structure_kernel import ProductStructureKernel
 from .rbf_kernel import RBFKernel
 from .scale_kernel import ScaleKernel
 from .spectral_mixture_kernel import SpectralMixtureKernel
@@ -24,7 +24,7 @@ from .white_noise_kernel import WhiteNoiseKernel
 __all__ = [
     "Kernel",
     "AdditiveKernel",
-    "AdditiveGridInterpolationKernel",
+    "AdditiveStructureKernel",
     "CosineKernel",
     "GridKernel",
     "GridInterpolationKernel",
@@ -33,10 +33,10 @@ __all__ = [
     "LCMKernel",
     "LinearKernel",
     "MaternKernel",
-    "MultiplicativeGridInterpolationKernel",
     "MultitaskKernel",
     "PeriodicKernel",
     "ProductKernel",
+    "ProductStructureKernel",
     "RBFKernel",
     "ScaleKernel",
     "SpectralMixtureKernel",
diff --git a/gpytorch/kernels/additive_grid_interpolation_kernel.py b/gpytorch/kernels/additive_grid_interpolation_kernel.py
deleted file mode 100644
index c25c980d..00000000
--- a/gpytorch/kernels/additive_grid_interpolation_kernel.py
+++ /dev/null
@@ -1,99 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from __future__ import unicode_literals
-
-import warnings
-from .grid_interpolation_kernel import GridInterpolationKernel
-from ..utils.interpolation import Interpolation
-
-
-class AdditiveGridInterpolationKernel(GridInterpolationKernel):
-    r"""
-    A variant of :class:`~gpytorch.kernels.GridInterpolationKernel` designed specifically
-    for additive kernels. If a kernel decomposes additively, then this module will be much more
-    computationally efficient.
-
-    A kernel function `k` decomposes additively if it can be written as
-
-    .. math::
-
-       \begin{equation*}
-          k(\mathbf{x_1}, \mathbf{x_2}) = k'(x_1^{(1)}, x_2^{(1)}) + \ldots + k'(x_1^{(d)}, x_2^{(d)})
-       \end{equation*}
-
-    for some kernel :math:`k'` that operates on a subset of dimensions.
-
-    The groupings of dimensions are specified by the :attr:`dim_groups` attribute.
-
-    * `dim_groups=d` (d is the dimensionality of :math:`\mathbf x`):
-        the kernel :math:`k` will be the sum of `d` sub-kernels, each operating
-        on one dimension of :math:`\mathbf x`.
-
-    * `dim_groups=d/2`:
-        the first sub-kernel operates on dimensions 1 and 2, the second sub-kernel
-        operates on dimensions 3 and 4, etc.
-
-    * `dim_groups=1`:
-        there is no additive decomposition
-
-    .. note::
-
-        `AdditiveGridInterpolationKernel` can only wrap **stationary kernels** (such as RBF, Matern,
-        Periodic, Spectral Mixture, etc.)
-
-    Args:
-        :attr:`base_kernel_module` (Kernel):
-            The kernel to approximate with KISS-GP
-        :attr:`grid_size` (int):
-            The size of the grid (in each dimension)
-        :attr:`num_dims` (int):
-            The dimension of the input data. Required if `grid_bounds=None`
-        :attr:`dim_groups` (int):
-            The number of additive components
-        :attr:`grid_bounds` (tuple(float, float), optional):
-            The bounds of the grid, if known (high performance mode).
-            The length of the tuple must match the size of the dim group (`num_dims // dim_groups`).
-            The entries represent the min/max values for each dimension.
-        :attr:`active_dims` (tuple of ints, optional):
-            Passed down to the `base_kernel_module`.
-    """
-
-    def __init__(
-        self,
-        base_kernel_module,
-        grid_size,
-        dim_groups=None,
-        num_dims=None,
-        grid_bounds=None,
-        active_dims=None,
-        n_components=None,
-    ):
-        if n_components is not None:
-            warnings.warn("n_components is deprecated. Use dim_groups instead.", DeprecationWarning)
-            dim_groups = n_components
-        if dim_groups is None:
-            raise RuntimeError("Must supply dim_groups")
-
-        super(AdditiveGridInterpolationKernel, self).__init__(
-            base_kernel_module, grid_size, num_dims // dim_groups, grid_bounds, active_dims=active_dims
-        )
-
-        self.dim_groups = dim_groups
-
-    def _compute_grid(self, inputs):
-        inputs = inputs.view(inputs.size(0), inputs.size(1), self.dim_groups, -1)
-        batch_size, n_data, dim_groups, n_dimensions = inputs.size()
-        inputs = inputs.transpose(0, 2).contiguous().view(dim_groups * batch_size * n_data, n_dimensions)
-        interp_indices, interp_values = Interpolation().interpolate(self.grid, inputs)
-        interp_indices = interp_indices.view(dim_groups * batch_size, n_data, -1)
-        interp_values = interp_values.view(dim_groups * batch_size, n_data, -1)
-        return interp_indices, interp_values
-
-    def _inducing_forward(self):
-        res = super(AdditiveGridInterpolationKernel, self)._inducing_forward()
-        return res.repeat(self.dim_groups, 1, 1)
-
-    def forward(self, x1, x2):
-        res = super(AdditiveGridInterpolationKernel, self).forward(x1, x2)
-        return res.sum_batch(sum_batch_size=self.dim_groups)
diff --git a/gpytorch/kernels/additive_structure_kernel.py b/gpytorch/kernels/additive_structure_kernel.py
new file mode 100644
index 00000000..ff215d90
--- /dev/null
+++ b/gpytorch/kernels/additive_structure_kernel.py
@@ -0,0 +1,61 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from .kernel import Kernel
+from ..lazy import LazyTensor, NonLazyTensor
+
+
+class AdditiveStructureKernel(Kernel):
+    r"""
+    A Kernel decorator for kernels with additive structure. If a kernel decomposes
+    additively, then this module will be much more computationally efficient.
+
+    A kernel function `k` decomposes additively if it can be written as
+
+    .. math::
+
+       \begin{equation*}
+          k(\mathbf{x_1}, \mathbf{x_2}) = k'(x_1^{(1)}, x_2^{(1)}) + \ldots + k'(x_1^{(d)}, x_2^{(d)})
+       \end{equation*}
+
+    for some kernel :math:`k'` that operates on a subset of dimensions.
+
+    Given a `b x n x d` input, `AdditiveStructureKernel` computes `d` one-dimensional kernels
+    (using the supplied base_kernel), and then adds the component kernels together.
+    Unlike :class:`~gpytorch.kernels.AdditiveKernel`, `AdditiveStructureKernel` computes each
+    of the additive terms in batch, making it very fast.
+
+    Args:
+        :attr:`base_kernel` (Kernel):
+            The kernel to approximate with KISS-GP
+        :attr:`num_dims` (int):
+            The dimension of the input data.
+        :attr:`active_dims` (tuple of ints, optional):
+            Passed down to the `base_kernel`.
+    """
+
+    def __init__(self, base_kernel, num_dims, active_dims=None):
+        super(AdditiveStructureKernel, self).__init__(active_dims=active_dims)
+        self.base_kernel = base_kernel
+        self.num_dims = num_dims
+
+    def forward(self, x1, x2, batch_dims=None, **params):
+        if batch_dims == (0, 2):
+            raise RuntimeError(
+                "AdditiveStructureKernel does not accept the batch_dims argument."
+            )
+
+        res = self.base_kernel(x1, x2, batch_dims=(0, 2), **params).evaluate_kernel()
+
+        evaluate = False
+        if not isinstance(res, LazyTensor):
+            evaluate = True
+            res = NonLazyTensor(res)
+
+        res = res.sum_batch(sum_batch_size=x1.size(-1))
+
+        if evaluate:
+            res = res.evaluate()
+        return res
diff --git a/gpytorch/kernels/cosine_kernel.py b/gpytorch/kernels/cosine_kernel.py
index c7b24dab..03536279 100644
--- a/gpytorch/kernels/cosine_kernel.py
+++ b/gpytorch/kernels/cosine_kernel.py
@@ -65,10 +65,10 @@ class CosineKernel(Kernel):
     def period_length(self):
         return self.log_period_length.exp().clamp(self.eps, 1e5)
 
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, **params):
         x1_ = x1.div(self.period_length)
         x2_ = x2.div(self.period_length)
-        x1_, x2_ = self._create_input_grid(x1_, x2_)
+        x1_, x2_ = self._create_input_grid(x1_, x2_, **params)
 
         diff = torch.norm((x1_ - x2_).abs(), 2, -1)
         res = torch.cos(diff.mul(math.pi))
diff --git a/gpytorch/kernels/grid_interpolation_kernel.py b/gpytorch/kernels/grid_interpolation_kernel.py
index bb0245a5..9dffff46 100644
--- a/gpytorch/kernels/grid_interpolation_kernel.py
+++ b/gpytorch/kernels/grid_interpolation_kernel.py
@@ -4,7 +4,6 @@ from __future__ import print_function
 from __future__ import unicode_literals
 
 import torch
-from .kernel import Kernel
 from .grid_kernel import GridKernel
 from ..lazy import InterpolatedLazyTensor
 from ..utils.interpolation import Interpolation
@@ -48,7 +47,7 @@ class GridInterpolationKernel(GridKernel):
         Periodic, Spectral Mixture, etc.)
 
     Args:
-        :attr:`base_kernel_module` (Kernel):
+        :attr:`base_kernel` (Kernel):
             The kernel to approximate with KISS-GP
         :attr:`grid_size` (int):
             The size of the grid (in each dimension)
@@ -59,13 +58,13 @@ class GridInterpolationKernel(GridKernel):
             The length of the tuple must match the number of dimensions.
             The entries represent the min/max values for each dimension.
         :attr:`active_dims` (tuple of ints, optional):
-            Passed down to the `base_kernel_module`.
+            Passed down to the `base_kernel`.
 
     .. Kernel Interpolation for Scalable Structured Gaussian Processes:
         http://proceedings.mlr.press/v37/wilson15.pdf
     """
 
-    def __init__(self, base_kernel_module, grid_size, num_dims=None, grid_bounds=None, active_dims=None):
+    def __init__(self, base_kernel, grid_size, num_dims=None, grid_bounds=None, active_dims=None):
         has_initialized_grid = 0
         grid_is_dynamic = True
 
@@ -95,15 +94,15 @@ class GridInterpolationKernel(GridKernel):
         inducing_points, grid = self._create_grid()
 
         super(GridInterpolationKernel, self).__init__(
-            base_kernel_module=base_kernel_module, inducing_points=inducing_points, grid=grid, active_dims=active_dims
+            base_kernel=base_kernel, inducing_points=inducing_points, grid=grid, active_dims=active_dims
         )
         self.register_buffer("has_initialized_grid", torch.tensor(has_initialized_grid, dtype=torch.uint8))
 
     def _create_grid(self):
-        grid = torch.zeros(len(self.grid_bounds), self.grid_size)
+        grid = torch.zeros(self.grid_size, len(self.grid_bounds))
         for i in range(len(self.grid_bounds)):
             grid_diff = float(self.grid_bounds[i][1] - self.grid_bounds[i][0]) / (self.grid_size - 2)
-            grid[i] = torch.linspace(
+            grid[:, i] = torch.linspace(
                 self.grid_bounds[i][0] - grid_diff, self.grid_bounds[i][1] + grid_diff, self.grid_size
             )
 
@@ -111,7 +110,7 @@ class GridInterpolationKernel(GridKernel):
         prev_points = None
         for i in range(len(self.grid_bounds)):
             for j in range(self.grid_size):
-                inducing_points[j * self.grid_size ** i : (j + 1) * self.grid_size ** i, i].fill_(grid[i, j])
+                inducing_points[j * self.grid_size ** i : (j + 1) * self.grid_size ** i, i].fill_(grid[j, i])
                 if prev_points is not None:
                     inducing_points[j * self.grid_size ** i : (j + 1) * self.grid_size ** i, :i].copy_(prev_points)
             prev_points = inducing_points[: self.grid_size ** (i + 1), : (i + 1)]
@@ -130,21 +129,26 @@ class GridInterpolationKernel(GridKernel):
     def has_custom_exact_predictions(self):
         return True
 
-    def _compute_grid(self, inputs):
+    def _compute_grid(self, inputs, batch_dims):
         batch_size, n_data, n_dimensions = inputs.size()
+        if batch_dims == (0, 2):
+            inputs = inputs.view(inputs.size(0), inputs.size(1), -1, 1)
+            inputs = inputs.transpose(1, 2).contiguous()
+            batch_size = batch_size * inputs.size(1)
+            n_dimensions = n_dimensions // inputs.size(1)
+
         inputs = inputs.view(batch_size * n_data, n_dimensions)
         interp_indices, interp_values = Interpolation().interpolate(self.grid, inputs)
         interp_indices = interp_indices.view(batch_size, n_data, -1)
         interp_values = interp_values.view(batch_size, n_data, -1)
         return interp_indices, interp_values
 
-    def _inducing_forward(self):
-        return super(GridInterpolationKernel, self).forward(self.inducing_points, self.inducing_points)
-
-    def forward_diag(self, x1, x2, **kwargs):
-        return super(Kernel, self).__call__(x1, x2, **kwargs).diag().unsqueeze(-1)
+    def _inducing_forward(self, batch_dims, **params):
+        return super(GridInterpolationKernel, self).forward(
+            self.inducing_points, self.inducing_points, batch_dims=batch_dims, **params,
+        )
 
-    def forward(self, x1, x2, **kwargs):
+    def forward(self, x1, x2, batch_dims=None, **params):
         # See if we need to update the grid or not
         if self.grid_is_dynamic:  # This is true if a grid_bounds wasn't passed in
             if torch.equal(x1, x2):
@@ -172,20 +176,23 @@ class GridInterpolationKernel(GridKernel):
                 inducing_points, grid = self._create_grid()
                 self.update_inducing_points_and_grid(inducing_points, grid)
 
-        base_lazy_tsr = self._inducing_forward()
+        base_lazy_tsr = self._inducing_forward(batch_dims=batch_dims, **params)
         if x1.size(0) > 1:
             base_lazy_tsr = base_lazy_tsr.repeat(x1.size(0), 1, 1)
 
-        left_interp_indices, left_interp_values = self._compute_grid(x1)
+        left_interp_indices, left_interp_values = self._compute_grid(x1, batch_dims)
         if torch.equal(x1, x2):
             right_interp_indices = left_interp_indices
             right_interp_values = left_interp_values
         else:
-            right_interp_indices, right_interp_values = self._compute_grid(x2)
-        return InterpolatedLazyTensor(
+            right_interp_indices, right_interp_values = self._compute_grid(x2, batch_dims)
+
+        res = InterpolatedLazyTensor(
             base_lazy_tsr,
             left_interp_indices.detach(),
             left_interp_values,
             right_interp_indices.detach(),
             right_interp_values,
         )
+
+        return res
diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index fd21598d..1dce7b59 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -23,22 +23,22 @@ class GridKernel(Kernel):
         Periodic, Spectral Mixture, etc.)
 
     Args:
-        :attr:`base_kernel_module` (Kernel):
+        :attr:`base_kernel` (Kernel):
             The kernel to speed up with grid methods.
         :attr:`inducing_points` (Tensor, n x d):
             This will be the set of points that lie on the grid.
         :attr:`grid` (Tensor, k x d):
             The exact grid points.
         :attr:`active_dims` (tuple of ints, optional):
-            Passed down to the `base_kernel_module`.
+            Passed down to the `base_kernel`.
 
     .. _Fast kernel learning for multidimensional pattern extrapolation:
         http://www.cs.cmu.edu/~andrewgw/manet.pdf
     """
 
-    def __init__(self, base_kernel_module, inducing_points, grid, active_dims=None):
+    def __init__(self, base_kernel, inducing_points, grid, active_dims=None):
         super(GridKernel, self).__init__(active_dims=active_dims)
-        self.base_kernel_module = base_kernel_module
+        self.base_kernel = base_kernel
         if inducing_points.ndimension() != 2:
             raise RuntimeError("Inducing points should be 2 dimensional")
         self.register_buffer("inducing_points", inducing_points.unsqueeze(0))
@@ -59,27 +59,33 @@ class GridKernel(Kernel):
             del self._cached_kernel_mat
         return self
 
-    def forward(self, x1, x2, **kwargs):
+    def forward(self, x1, x2, diag=False, batch_dims=None, **params):
         if not torch.equal(x1, self.inducing_points) or not torch.equal(x2, self.inducing_points):
             raise RuntimeError("The kernel should only receive the inducing points as input")
 
         if not self.training and hasattr(self, "_cached_kernel_mat"):
-            return self._cached_kernel_mat
+            covar = self._cached_kernel_mat
 
         else:
             n_dim = x1.size(-1)
-            grid_var = self.grid.view(n_dim, -1, 1)
+            grid = self.grid.unsqueeze(0)
 
             if settings.use_toeplitz.on():
-                first_item = grid_var[:, 0:1].contiguous()
-                covar_columns = self.base_kernel_module(first_item, grid_var, **kwargs).evaluate()
-                covars = [ToeplitzLazyTensor(covar_columns[i : i + 1].squeeze(-2)) for i in range(n_dim)]
+                first_item = grid[:, 0:1]
+                covar_columns = self.base_kernel(first_item, grid, diag=False, batch_dims=(0, 2), **params)
+                covar_columns = covar_columns.evaluate().squeeze(-2)
+                if batch_dims == (0, 2):
+                    covars = [ToeplitzLazyTensor(covar_columns)]
+                else:
+                    covars = [ToeplitzLazyTensor(covar_columns[i : i + 1]) for i in range(n_dim)]
             else:
-                grid_var = grid_var.view(n_dim, -1, 1)
-                covars = self.base_kernel_module(grid_var, grid_var, **kwargs).evaluate_kernel()
-                covars = [covars[i : i + 1] for i in range(n_dim)]
+                covars = self.base_kernel(grid, grid, batch_dims=(0, 2), **params).evaluate_kernel()
+                if batch_dims == (0, 2):
+                    covars = [covars]
+                else:
+                    covars = [covars[i : i + 1] for i in range(n_dim)]
 
-            if n_dim > 1:
+            if len(covars) > 1:
                 covar = KroneckerProductLazyTensor(*covars)
             else:
                 covar = covars[0]
@@ -87,4 +93,4 @@ class GridKernel(Kernel):
             if not self.training:
                 self._cached_kernel_mat = covar
 
-            return covar
+        return covar
diff --git a/gpytorch/kernels/index_kernel.py b/gpytorch/kernels/index_kernel.py
index 370c926f..cd763318 100644
--- a/gpytorch/kernels/index_kernel.py
+++ b/gpytorch/kernels/index_kernel.py
@@ -64,7 +64,7 @@ class IndexKernel(Kernel):
         res = PsdSumLazyTensor(RootLazyTensor(self.covar_factor), DiagLazyTensor(self.log_var.exp()))
         return res
 
-    def forward(self, i1, i2):
+    def forward(self, i1, i2, **params):
         covar_matrix = _eval_covar_matrix(self.covar_factor, self.log_var)
         res = InterpolatedLazyTensor(base_lazy_tensor=covar_matrix, left_interp_indices=i1, right_interp_indices=i2)
         return res
diff --git a/gpytorch/kernels/inducing_point_kernel.py b/gpytorch/kernels/inducing_point_kernel.py
index bc3708bb..a7c4287e 100644
--- a/gpytorch/kernels/inducing_point_kernel.py
+++ b/gpytorch/kernels/inducing_point_kernel.py
@@ -12,9 +12,9 @@ from ..variational import MVNVariationalStrategy
 
 
 class InducingPointKernel(Kernel):
-    def __init__(self, base_kernel_module, inducing_points, active_dims=None):
+    def __init__(self, base_kernel, inducing_points, active_dims=None):
         super(InducingPointKernel, self).__init__(active_dims=active_dims)
-        self.base_kernel_module = base_kernel_module
+        self.base_kernel = base_kernel
 
         if inducing_points.ndimension() == 1:
             inducing_points = inducing_points.unsqueeze(-1)
@@ -33,7 +33,7 @@ class InducingPointKernel(Kernel):
         if not self.training and hasattr(self, "_cached_kernel_mat"):
             return self._cached_kernel_mat
         else:
-            res = self.base_kernel_module(self.inducing_points, self.inducing_points).evaluate()
+            res = self.base_kernel(self.inducing_points, self.inducing_points).evaluate()
             if not self.training:
                 self._cached_kernel_mat = res
             return res
@@ -56,11 +56,11 @@ class InducingPointKernel(Kernel):
             return res
 
     def _get_covariance(self, x1, x2):
-        k_ux1 = self.base_kernel_module(x1, self.inducing_points).evaluate()
+        k_ux1 = self.base_kernel(x1, self.inducing_points).evaluate()
         if torch.equal(x1, x2):
             covar = RootLazyTensor(k_ux1.matmul(self._inducing_inv_root))
         else:
-            k_ux2 = self.base_kernel_module(x2, self.inducing_points).evaluate()
+            k_ux2 = self.base_kernel(x2, self.inducing_points).evaluate()
             covar = MatmulLazyTensor(
                 k_ux1.matmul(self._inducing_inv_root), k_ux2.matmul(self._inducing_inv_root).transpose(-1, -2)
             )
@@ -75,7 +75,7 @@ class InducingPointKernel(Kernel):
         inputs = inputs.unsqueeze(-2).view(-1, 1, inputs.size(-1))
 
         # Get diagonal of covar
-        covar_diag = self.base_kernel_module(inputs)
+        covar_diag = self.base_kernel(inputs)
         if isinstance(covar_diag, LazyTensor):
             covar_diag = covar_diag.evaluate()
         covar_diag = covar_diag.view(orig_size[:-1])
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 27b04821..c4e85659 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -6,9 +6,9 @@ from __future__ import unicode_literals
 from abc import abstractmethod
 import torch
 from torch.nn import ModuleList
-from ..lazy import LazyTensor, LazyEvaluatedKernelTensor, ZeroLazyTensor
+from ..lazy import LazyEvaluatedKernelTensor, ZeroLazyTensor
 from ..module import Module
-from ..utils import prod
+from .. import settings
 
 
 class Kernel(Module):
@@ -131,15 +131,7 @@ class Kernel(Module):
             return torch.Size(non_batch_size)
 
     @abstractmethod
-    def forward_diag(self, x1, x2, **params):
-        res = super(Kernel, self).__call__(x1.transpose(-2, -3), x2.transpose(-2, -3), **params)
-        if isinstance(res, LazyTensor):
-            res = res.evaluate()
-        res = res.transpose(-2, -3)
-        return res
-
-    @abstractmethod
-    def forward(self, x1, x2, **params):
+    def forward(self, x1, x2, diag=False, batch_dims=None, **params):
         """
         Computes the covariance between x1 and x2.
         This method should be imlemented by all Kernel subclasses.
@@ -154,20 +146,28 @@ class Kernel(Module):
 
         Args:
             :attr:`x1` (Tensor `n x d` or `b x n x d`)
-            :attr:`x2` (Tensor `m x d` or `b x m x d`) - for diag mode, these must be the same inputs
+            :attr:`x2` (Tensor `m x d` or `b x m x d`)
+            :attr:`diag` (bool):
+                Should the Kernel compute the whole kernel, or just the diag?
+                For most Kernels, this option will be passed into `create_input_grid`
+            :attr:`batch_dims` (tuple, optional):
+                If this option is passed in, it will tell the tensor which of the
+                three dimensions are batch dimensions.
+                Currently accepts: standard mode (either None or (0,))
+                or (0, 2) for use with Additive/Multiplicative kernels
 
         Returns:
             :class:`Tensor` or :class:`gpytorch.lazy.LazyTensor`.
             The exact size depends on the kernel's evaluation mode:
 
             * `full_covar`: `n x m` or `b x n x m`
-            * `full_covar` with `dim_groups=k`: `k x n x m` or `b x k x n x m`
+            * `full_covar` with `batch_dims=(0, 2)`: `k x n x m` or `b x k x n x m`
             * `diag`: `n` or `b x n`
-            * `diag` with `dim_groups=k`: `k x n` or `b x k x n`
+            * `diag` with `batch_dims=(0, 2)`: `k x n` or `b x k x n`
         """
         raise NotImplementedError()
 
-    def _create_input_grid(self, x1, x2):
+    def _create_input_grid(self, x1, x2, diag=False, batch_dims=None, **params):
         """
         This is a helper method for creating a grid of the kernel's inputs.
         Use this helper rather than maually creating a meshgrid.
@@ -183,13 +183,28 @@ class Kernel(Module):
             The shape depends on the kernel's mode
 
             * `full_covar`: (`b x n x 1 x d` and `b x 1 x m x d`)
-            * `full_covar` with `dim_groups=k`: (`b x k x n x 1 x (d//k)` and `b x k x 1 x m x (d//k)`)
+            * `full_covar` with `batch_dims=(0, 2)`: (`b x k x n x 1 x 1` and `b x k x 1 x m x 1`)
             * `diag`: (`b x n x d` and `b x n x d`)
-            * `diag` with `dim_groups=k`: (`b x k x n x (d//k)` and `b x k x n x (d//k)`)
+            * `diag` with `batch_dims=(0, 2)`: (`b x k x n x 1` and `b x k x n x 1`)
         """
-        return x1.unsqueeze(-2), x2.unsqueeze(-3)
+        x1_, x2_ = x1, x2
+        if batch_dims == (0, 2):
+            x1_ = x1_.view(*x1.size()[:-1], -1, 1)
+            x1_ = x1_.permute(0, -2, *list(range(1, x1_.dim() - 2)), -1).contiguous()
+            x1_ = x1_.view(-1, *x1_.size()[2:])
+            if torch.equal(x1, x2):
+                x2_ = x1_
+            else:
+                x2_ = x2_.view(*x2.size()[:-1], -1, 1)
+                x2_ = x2_.permute(0, -2, *list(range(1, x2_.dim() - 2)), -1).contiguous()
+                x2_ = x2_.view(-1, *x2_.size()[2:])
+
+        if diag:
+            return x1_, x2_
+        else:
+            return x1_.unsqueeze(-2), x2_.unsqueeze(-3)
 
-    def __call__(self, x1, x2=None, **params):
+    def __call__(self, x1, x2=None, diag=False, batch_dims=None, **params):
         x1_, x2_ = x1, x2
 
         # Select the active dimensions
@@ -210,7 +225,78 @@ class Kernel(Module):
         if x2_ is None:
             x2_ = x1_
 
-        return LazyEvaluatedKernelTensor(self, x1_, x2_, **params)
+        # Check batch_dims args
+        if settings.debug.on():
+            if batch_dims is not None:
+                if not isinstance(batch_dims, tuple) or (batch_dims != (0,) and batch_dims != (0, 2)):
+                    raise RuntimeError(
+                        "batch_dims currently accepts either None, (0,), or (0, 2). Got {}.".format(batch_dims)
+                    )
+
+        if diag:
+            res = super(Kernel, self).__call__(x1_, x2_, diag=True, batch_dims=batch_dims, **params)
+
+            # Did this Kernel eat the diag option?
+            # If it does not return a LazyEvaluatedKernelTensor, we can call diag on the output
+            if not isinstance(res, LazyEvaluatedKernelTensor):
+                if res.dim() == x1_.dim() and res.shape[-2:] == torch.Size((x1_.size(-2), x2_.size(-2))):
+                    res = res.diag()
+
+            # Now we'll make sure that the shape we're getting from diag makes sense
+            if settings.debug.on():
+                # If we used batch_dims...
+                shape = self.size(x1_, x2_)
+                if batch_dims == (0, 2):
+                    if len(shape) == 2:
+                        expected_shape = torch.Size((x1_.size(-1), shape[0]))
+                    else:
+                        expected_shape = torch.Size((shape[0] * x1_.size(-1), shape[1]))
+                    if res.shape != expected_shape:
+                        raise RuntimeError(
+                            "The kernel {} is not equipped to handle batch_dims=(0, 2) "
+                            "and diag. Expected size {}. Got size {}.".format(
+                                self.__class__.__name__, expected_shape, res.shape
+                            )
+                        )
+
+                # If we didn't use batch_dims...
+                else:
+                    expected_shape = shape[:-1]
+                    if res.shape != expected_shape:
+                        raise RuntimeError(
+                            "The kernel {} is not equipped to handle and diag. Expected size {}. "
+                            "Got size {}".format(self.__class__.__name__, expected_shape, res.shape)
+                        )
+            return res
+
+        else:
+            res = LazyEvaluatedKernelTensor(self, x1_, x2_, batch_dims=batch_dims, **params)
+
+            # Now we'll make sure that the shape we're getting makes sense
+            if settings.debug.on():
+                # If we used batch_dims...
+                shape = self.size(x1_, x2_)
+                if batch_dims == (0, 2):
+                    if len(shape) == 2:
+                        expected_shape = torch.Size((x1_.size(-1), shape[0], shape[1]))
+                    else:
+                        expected_shape = torch.Size((shape[0] * x1_.size(-1), shape[1], shape[2]))
+                    if res.shape != expected_shape:
+                        raise RuntimeError(
+                            "The kernel {} is not equipped to handle batch_dims=(0, 2). Expected size {}. "
+                            "Got size {}".format(self.__class__.__name__, expected_shape, res.shape)
+                        )
+
+                # If we didn't use batch_dims...
+                else:
+                    expected_shape = shape
+                    if res.shape != expected_shape:
+                        raise RuntimeError(
+                            "Error with {}.forward. Expected size {}. Got size {}.".format(
+                                self.__class__.__name__, expected_shape, res.shape
+                            )
+                        )
+            return res
 
     def __add__(self, other):
         return AdditiveKernel(self, other)
@@ -233,10 +319,13 @@ class AdditiveKernel(Kernel):
         super(AdditiveKernel, self).__init__()
         self.kernels = ModuleList(kernels)
 
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, **params):
         res = ZeroLazyTensor()
         for kern in self.kernels:
-            res = res + kern(x1, x2).evaluate_kernel()
+            next_term = kern(x1, x2, **params)
+            if isinstance(next_term, LazyEvaluatedKernelTensor):
+                next_term = next_term.evaluate_kernel()
+            res = res + next_term
         return res
 
 
@@ -254,5 +343,11 @@ class ProductKernel(Kernel):
         super(ProductKernel, self).__init__()
         self.kernels = ModuleList(kernels)
 
-    def forward(self, x1, x2):
-        return prod([k(x1, x2).evaluate_kernel() for k in self.kernels])
+    def forward(self, x1, x2, **params):
+        res = self.kernels[0](x1, x2, **params)
+        for kern in self.kernels[1:]:
+            next_term = kern(x1, x2, **params)
+            if isinstance(next_term, LazyEvaluatedKernelTensor):
+                next_term = next_term.evaluate_kernel()
+            res = res * next_term
+        return res
diff --git a/gpytorch/kernels/lcm_kernel.py b/gpytorch/kernels/lcm_kernel.py
index 1a00fcbc..2ccb6a09 100644
--- a/gpytorch/kernels/lcm_kernel.py
+++ b/gpytorch/kernels/lcm_kernel.py
@@ -38,22 +38,10 @@ class LCMKernel(Kernel):
             for base_kernel in base_kernels
         ])
 
-    def forward_diag(self, x1, x2):
-        """
-        Returns the diagonal of the covariance matrix only. This overrides the
-        default behavior for this supplied in :class:`gpytorch.kernels.Kernel`
-        because we need to take the Kronecker product of the diagonals of the
-        base data kernel and the task kernel, and add them up.
-        """
-        res = self.covar_module_list[0].forward_diag(x1, x2)
-        for m in self.covar_module_list[1:]:
-            res += m.forward_diag(x1, x2)
-        return res
-
-    def forward(self, x1, x2):
-        res = self.covar_module_list[0].forward(x1, x2)
+    def forward(self, x1, x2, **params):
+        res = self.covar_module_list[0].forward(x1, x2, **params)
         for m in self.covar_module_list[1:]:
-            res += m.forward(x1, x2)
+            res += m.forward(x1, x2, **params)
         return res
 
     def size(self, x1, x2):
diff --git a/gpytorch/kernels/linear_kernel.py b/gpytorch/kernels/linear_kernel.py
index f3b77760..44f6e99d 100644
--- a/gpytorch/kernels/linear_kernel.py
+++ b/gpytorch/kernels/linear_kernel.py
@@ -47,25 +47,32 @@ class LinearKernel(Kernel):
             `len(active_dims)` should equal `num_dimensions`.
     """
 
-    def __init__(
-        self,
-        num_dimensions,
-        variance_prior=None,
-        offset_prior=None,
-        active_dims=None,
-    ):
+    def __init__(self, num_dimensions, variance_prior=None, offset_prior=None, active_dims=None):
         super(LinearKernel, self).__init__(active_dims=active_dims)
         self.register_parameter(name="variance", parameter=torch.nn.Parameter(torch.zeros(1)), prior=variance_prior)
         self.register_parameter(
             name="offset", parameter=torch.nn.Parameter(torch.zeros(1, 1, num_dimensions)), prior=offset_prior
         )
 
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, batch_dims=None, **params):
+        x1_ = x1 - self.offset
+        if batch_dims == (0, 2):
+            x1_ = x1_.view(x1_.size(0), x1_.size(1), -1, 1)
+            x1_ = x1_.permute(0, 2, 1, 3).contiguous()
+            x1_ = x1_.view(-1, x1_.size(-2), x1_.size(-1))
+
         if x1.size() == x2.size() and torch.equal(x1, x2):
             # Use RootLazyTensor when x1 == x2 for efficiency when composing
             # with other kernels
-            prod = RootLazyTensor(x1 - self.offset)
+            prod = RootLazyTensor(x1_)
+
         else:
-            prod = MatmulLazyTensor(x1 - self.offset, (x2 - self.offset).transpose(2, 1))
+            x2_ = x2 - self.offset
+            if batch_dims == (0, 2):
+                x2_ = x2_.view(x2_.size(0), x2_.size(1), -1, 1)
+                x2_ = x2_.permute(0, 2, 1, 3).contiguous()
+                x2_ = x2_.view(-1, x2_.size(-2), x2_.size(-1))
+
+            prod = MatmulLazyTensor(x1_, x2_.transpose(2, 1))
 
         return prod + self.variance.expand(prod.size())
diff --git a/gpytorch/kernels/matern_kernel.py b/gpytorch/kernels/matern_kernel.py
index 253eee4e..16283c7c 100644
--- a/gpytorch/kernels/matern_kernel.py
+++ b/gpytorch/kernels/matern_kernel.py
@@ -98,30 +98,12 @@ class MaternKernel(Kernel):
         )
         self.nu = nu
 
-    def forward_diag(self, x1, x2):
-        mean = x1.mean(1, keepdim=True).mean(0, keepdim=True)
-        x1_normed = x1 - mean.unsqueeze(0).unsqueeze(1)
-        x2_normed = x2 - mean.unsqueeze(0).unsqueeze(1)
-
-        diff = x1_normed - x2_normed
-        distance_over_rho = diff.pow_(2).sum(-1).sqrt()
-        exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance_over_rho)
-
-        if self.nu == 0.5:
-            constant_component = 1
-        elif self.nu == 1.5:
-            constant_component = (math.sqrt(3) * distance_over_rho).add(1)
-        elif self.nu == 2.5:
-            constant_component = (math.sqrt(5) * distance_over_rho).add(1).add(5. / 3. * distance_over_rho ** 2)
-
-        return constant_component * exp_component
-
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, **params):
         mean = x1.view(-1, 1, x1.size(-1)).mean(0, keepdim=True)
 
         x1_ = (x1 - mean).div(self.lengthscale)
         x2_ = (x2 - mean).div(self.lengthscale)
-        x1_, x2_ = self._create_input_grid(x1_, x2_)
+        x1_, x2_ = self._create_input_grid(x1_, x2_, **params)
 
         distance = (x1_ - x2_).norm(2, dim=-1)
         exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)
diff --git a/gpytorch/kernels/multiplicative_grid_interpolation_kernel.py b/gpytorch/kernels/multiplicative_grid_interpolation_kernel.py
deleted file mode 100644
index be14ed3f..00000000
--- a/gpytorch/kernels/multiplicative_grid_interpolation_kernel.py
+++ /dev/null
@@ -1,123 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-from __future__ import unicode_literals
-
-import warnings
-from .grid_interpolation_kernel import GridInterpolationKernel
-from ..utils.interpolation import Interpolation
-
-
-class MultiplicativeGridInterpolationKernel(GridInterpolationKernel):
-    r"""
-    A variant of :class:`~gpytorch.kernels.GridInterpolationKernel` designed specifically
-    for kernels with product structure. If a kernel decomposes
-    multiplicatively, then this module will be much more computationally efficient.
-    See `Product Kernel Interpolation for Scalable Gaussian Processes`_ for more detail.
-
-    A kernel function `k` has product structure if it can be written as
-
-    .. math::
-
-       \begin{equation*}
-          k(\mathbf{x_1}, \mathbf{x_2}) = k'(x_1^{(1)}, x_2^{(1)}) * \ldots * k'(x_1^{(d)}, x_2^{(d)})
-       \end{equation*}
-
-    for some kernel :math:`k'` that operates on a subset of dimensions.
-
-    The groupings of dimensions are specified by the :attr:`dim_groups` attribute.
-
-    * `dim_groups=d` (d is the dimensionality of :math:`\mathbf x`):
-        the kernel :math:`k` will be the sum of `d` sub-kernels, each operating on one dimension of :math:`\mathbf x`.
-    * `dim_groups=d/2`:
-        the first sub-kernel operates on dimensions 1 and 2, the second sub-kernel
-        operates on dimensions 3 and 4, etc.
-    * `dim_groups=1`:
-        there is no multiplicative decomposition
-
-    .. note::
-
-        `AdditiveGridInterpolationKernel` can only wrap **stationary kernels** (such as RBF, Matern,
-        Periodic, Spectral Mixture, etc.)
-
-    .. note::
-
-        The :class:`~gpytorch.kernels.RBFKernel` decomposes multiplicatively. You should use
-        `MultiplicativeGridInterpolationKernel` for multi-dimension RBF kernels!
-
-    Args:
-        :attr:`base_kernel_module` (Kernel):
-            The kernel to approximate with KISS-GP
-        :attr:`grid_size` (int):
-            The size of the grid (in each dimension)
-        :attr:`num_dims` (int):
-            The dimension of the input data. Required if `grid_bounds=None`
-        :attr:`dim_groups` (int):
-            The number of multiplicative components
-        :attr:`grid_bounds` (tuple(float, float), optional):
-            The bounds of the grid, if known (high performance mode).
-            The length of the tuple must match the size of the dim group (`num_dims // dim_groups`).
-            The entries represent the min/max values for each dimension.
-        :attr:`active_dims` (tuple of ints, optional):
-            Passed down to the `base_kernel_module`.
-
-    .. Product Kernel Interpolation for Scalable Gaussian Processes:
-        https://arxiv.org/pdf/1802.08903
-    """
-
-    def __init__(
-        self,
-        base_kernel_module,
-        grid_size,
-        dim_groups=None,
-        num_dims=None,
-        grid_bounds=None,
-        active_dims=None,
-        n_components=None,
-    ):
-        if n_components is not None:
-            warnings.warn("n_components is deprecated. Use dim_groups instead.", DeprecationWarning)
-            dim_groups = n_components
-        if dim_groups is None:
-            raise RuntimeError("Must supply dim_groups")
-
-        super(MultiplicativeGridInterpolationKernel, self).__init__(
-            base_kernel_module=base_kernel_module,
-            grid_size=grid_size,
-            num_dims=(num_dims // dim_groups),
-            grid_bounds=grid_bounds,
-            active_dims=active_dims,
-        )
-        self.dim_groups = dim_groups
-
-    def _compute_grid(self, inputs):
-        inputs = inputs.view(inputs.size(0), inputs.size(1), self.dim_groups, -1)
-        batch_size, n_data, dim_groups, n_dimensions = inputs.size()
-        inputs = inputs.transpose(0, 2).contiguous().view(dim_groups * batch_size * n_data, n_dimensions)
-        interp_indices, interp_values = Interpolation().interpolate(self.grid, inputs)
-        interp_indices = interp_indices.view(dim_groups * batch_size, n_data, -1)
-        interp_values = interp_values.view(dim_groups * batch_size, n_data, -1)
-        return interp_indices, interp_values
-
-    def _inducing_forward(self):
-        res = super(MultiplicativeGridInterpolationKernel, self)._inducing_forward()
-        return res.repeat(self.dim_groups, 1, 1)
-
-    def forward(self, x1, x2):
-        res = super(MultiplicativeGridInterpolationKernel, self).forward(x1, x2)
-        res = res.mul_batch(mul_batch_size=self.dim_groups)
-        return res
-
-    def __call__(self, x1_, x2_=None, **params):
-        """
-        We cannot lazily evaluate actual kernel calls when using SKIP, because we
-        cannot root decompose rectangular matrices.
-
-        Because we slice in to the kernel during prediction to get the test x train
-        covar before calling evaluate_kernel, the order of operations would mean we
-        would get a MulLazyTensor representing a rectangular matrix, which we
-        cannot matmul with because we cannot root decompose it. Thus, SKIP actually
-        *requires* that we work with the full (train + test) x (train + test)
-        kernel matrix.
-        """
-        return super(MultiplicativeGridInterpolationKernel, self).__call__(x1_, x2_, **params).evaluate_kernel()
diff --git a/gpytorch/kernels/multitask_kernel.py b/gpytorch/kernels/multitask_kernel.py
index 299dba94..923c7d01 100644
--- a/gpytorch/kernels/multitask_kernel.py
+++ b/gpytorch/kernels/multitask_kernel.py
@@ -41,39 +41,23 @@ class MultitaskKernel(Kernel):
         self.n_tasks = n_tasks
         self.batch_size = 1
 
-    def forward_diag(self, x1, x2):
-        """
-        Returns the diagonal of the covariance matrix only. This overrides the default behavior for this supplied
-        in :class:`gpytorch.kernels.Kernel` because we need to take the Kronecker product of the diagonals of the
-        base data kernel and the task kernel.
-        """
-        task_indices = torch.arange(self.n_tasks, device=x1.device).long()
-        task_indices = task_indices.unsqueeze(0).unsqueeze(-1)
-
-        # These are small because they are vectors, therefore it is safe to evaluate them
-        covar_i_diag = self.task_covar_module.forward_diag(task_indices, task_indices)
-        covar_x_diag = self.data_covar_module.forward_diag(x1, x2)
-
-        if isinstance(covar_x_diag, LazyTensor):
-            covar_x_diag = covar_x_diag.evaluate()
-        if isinstance(covar_i_diag, LazyTensor):
-            covar_i_diag = covar_i_diag.evaluate()
+    def forward(self, x1, x2, diag=False, batch_dims=None, **params):
+        if batch_dims == (0, 2):
+            raise RuntimeError(
+                "AdditiveGridInterpolationKernel does not accept the batch_dims argument."
+            )
 
-        covar_x_diag = covar_x_diag.squeeze(-1)
-        covar_i_diag = covar_i_diag.squeeze(-1)
-
-        # Take the Kronecker product of the two diagonals
-        res = KroneckerProductLazyTensor(NonLazyTensor(covar_i_diag), NonLazyTensor(covar_x_diag)).evaluate()
-        return res.unsqueeze(-1)
-
-    def forward(self, x1, x2):
         covar_i = self.task_covar_module.covar_matrix
         covar_i = covar_i.repeat(x1.size(0), 1, 1)
-        covar_x = self.data_covar_module(x1, x2)
+        covar_x = self.data_covar_module(x1, x2, **params)
         if not isinstance(covar_x, LazyTensor):
             covar_x = NonLazyTensor(covar_x)
         res = KroneckerProductLazyTensor(covar_i, covar_x)
-        return res
+
+        if diag:
+            return res.diag()
+        else:
+            return res
 
     def size(self, x1, x2):
         """
diff --git a/gpytorch/kernels/periodic_kernel.py b/gpytorch/kernels/periodic_kernel.py
index afc08cba..72cff277 100644
--- a/gpytorch/kernels/periodic_kernel.py
+++ b/gpytorch/kernels/periodic_kernel.py
@@ -96,10 +96,10 @@ class PeriodicKernel(Kernel):
     def period_length(self):
         return self.log_period_length.exp().clamp(self.eps, 1e5)
 
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, **params):
         x1_ = x1.div(self.period_length)
         x2_ = x2.div(self.period_length)
-        x1_, x2_ = self._create_input_grid(x1_, x2_)
+        x1_, x2_ = self._create_input_grid(x1_, x2_, **params)
 
         diff = torch.sum((x1_ - x2_).abs(), -1)
         res = torch.sin(diff.mul(math.pi)).pow(2).mul(-2 / self.lengthscale).exp_()
diff --git a/gpytorch/kernels/product_structure_kernel.py b/gpytorch/kernels/product_structure_kernel.py
new file mode 100644
index 00000000..323401ac
--- /dev/null
+++ b/gpytorch/kernels/product_structure_kernel.py
@@ -0,0 +1,82 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+from .kernel import Kernel
+from ..lazy import LazyTensor, NonLazyTensor
+
+
+class ProductStructureKernel(Kernel):
+    r"""
+    A Kernel decorator for kernels with product structure. If a kernel decomposes
+    multiplicatively, then this module will be much more computationally efficient.
+
+    A kernel function `k` has product structure if it can be written as
+
+    .. math::
+
+       \begin{equation*}
+          k(\mathbf{x_1}, \mathbf{x_2}) = k'(x_1^{(1)}, x_2^{(1)}) * \ldots * k'(x_1^{(d)}, x_2^{(d)})
+       \end{equation*}
+
+    for some kernel :math:`k'` that operates on each dimension.
+
+    Given a `b x n x d` input, `ProductStructureKernel` computes `d` one-dimensional kernels
+    (using the supplied base_kernel), and then multiplies the component kernels together.
+    Unlike :class:`~gpytorch.kernels.ProductKernel`, `ProductStructureKernel` computes each
+    of the product terms in batch, making it very fast.
+
+    See `Product Kernel Interpolation for Scalable Gaussian Processes`_ for more detail.
+
+    Args:
+        :attr:`base_kernel` (Kernel):
+            The kernel to approximate with KISS-GP
+        :attr:`num_dims` (int):
+            The dimension of the input data.
+        :attr:`active_dims` (tuple of ints, optional):
+            Passed down to the `base_kernel`.
+
+    .. Product Kernel Interpolation for Scalable Gaussian Processes:
+        https://arxiv.org/pdf/1802.08903
+    """
+
+    def __init__(self, base_kernel, num_dims, active_dims=None):
+        super(ProductStructureKernel, self).__init__(active_dims=active_dims)
+        self.base_kernel = base_kernel
+        self.num_dims = num_dims
+
+    def forward(self, x1, x2, batch_dims=None, **params):
+        if batch_dims == (0, 2):
+            raise RuntimeError(
+                "ProductStructureKernel does not accept the batch_dims argument."
+            )
+
+        res = self.base_kernel(x1, x2, batch_dims=(0, 2), **params).evaluate_kernel()
+
+        evaluate = False
+        if not isinstance(res, LazyTensor):
+            evaluate = True
+            res = NonLazyTensor(res)
+
+        res = res.mul_batch(mul_batch_size=x1.size(-1))
+
+        if evaluate:
+            res = res.evaluate()
+        return res
+
+    def __call__(self, x1_, x2_=None, diag=False, batch_dims=None, **params):
+        """
+        We cannot lazily evaluate actual kernel calls when using SKIP, because we
+        cannot root decompose rectangular matrices.
+
+        Because we slice in to the kernel during prediction to get the test x train
+        covar before calling evaluate_kernel, the order of operations would mean we
+        would get a MulLazyTensor representing a rectangular matrix, which we
+        cannot matmul with because we cannot root decompose it. Thus, SKIP actually
+        *requires* that we work with the full (train + test) x (train + test)
+        kernel matrix.
+        """
+        return super(ProductStructureKernel, self).__call__(
+            x1_, x2_, diag=diag, batch_dims=batch_dims, **params
+        ).evaluate_kernel()
diff --git a/gpytorch/kernels/rbf_kernel.py b/gpytorch/kernels/rbf_kernel.py
index cd6b2120..2ae2b9ee 100644
--- a/gpytorch/kernels/rbf_kernel.py
+++ b/gpytorch/kernels/rbf_kernel.py
@@ -3,7 +3,6 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import math
 from .kernel import Kernel
 
 
@@ -66,14 +65,7 @@ class RBFKernel(Kernel):
         >>> covar = covar_module(x)  # Output: LazyTensor of size (2 x 10 x 10)
     """
 
-    def __init__(
-        self,
-        ard_num_dims=None,
-        log_lengthscale_prior=None,
-        eps=1e-6,
-        active_dims=None,
-        batch_size=1,
-    ):
+    def __init__(self, ard_num_dims=None, log_lengthscale_prior=None, eps=1e-6, active_dims=None, batch_size=1):
         super(RBFKernel, self).__init__(
             has_lengthscale=True,
             ard_num_dims=ard_num_dims,
@@ -83,16 +75,10 @@ class RBFKernel(Kernel):
             eps=eps,
         )
 
-    def forward_diag(self, x1, x2):
-        lengthscales = self.log_lengthscale.exp().mul(math.sqrt(2)).clamp(self.eps, 1e5)
-        diff = (x1 - x2).div_(lengthscales)
-        res = diff.pow_(2).sum(-1).mul_(-1).exp_().unsqueeze(-1)
-        return res
-
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, **params):
         x1_ = x1.div(self.lengthscale)
         x2_ = x2.div(self.lengthscale)
-        x1_, x2_ = self._create_input_grid(x1_, x2_)
+        x1_, x2_ = self._create_input_grid(x1_, x2_, **params)
 
         diff = (x1_ - x2_).norm(2, dim=-1)
         return diff.pow(2).div_(-2).exp_()
diff --git a/gpytorch/kernels/scale_kernel.py b/gpytorch/kernels/scale_kernel.py
index c75c862f..c3daa82b 100644
--- a/gpytorch/kernels/scale_kernel.py
+++ b/gpytorch/kernels/scale_kernel.py
@@ -58,12 +58,12 @@ class ScaleKernel(Kernel):
     def outputscale(self):
         return self.log_outputscale.exp()
 
-    def forward_diag(self, x1, x2):
+    def forward(self, x1, x2, batch_dims=None, **params):
         outputscales = self.log_outputscale.exp()
-        orig_output = self.base_kernel.forward_diag(x1, x2)
-        return orig_output.squeeze(-1).mul(outputscales.unsqueeze(-1))
+        if batch_dims == (0, 2) and outputscales.numel() > 1:
+            outputscales = outputscales.unsqueeze(1).repeat(1, x1.size(-1)).view(-1)
 
-    def forward(self, x1, x2):
-        outputscales = self.log_outputscale.exp()
-        orig_output = self.base_kernel(x1, x2)
+        orig_output = self.base_kernel(x1, x2, batch_dims=batch_dims, **params)
+        if torch.is_tensor(orig_output):
+            outputscales = outputscales.view(-1, *([1] * (orig_output.dim() - 1)))
         return orig_output.mul(outputscales)
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
index 3143d5af..c4088a22 100644
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -153,7 +153,7 @@ class SpectralMixtureKernel(Kernel):
         # Mixture weights should be roughly the stdv of the y values divided by the number of mixtures
         self.log_mixture_weights.data.fill_(train_y.std() / self.num_mixtures).log_()
 
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, **params):
         batch_size, n, num_dims = x1.size()
         _, m, _ = x2.size()
 
@@ -180,8 +180,8 @@ class SpectralMixtureKernel(Kernel):
         x2_cos = x2_ * self.mixture_means
 
         # Create grids
-        x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp)
-        x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos)
+        x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, **params)
+        x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, **params)
 
         # Compute the exponential and cosine terms
         exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi ** 2)
diff --git a/gpytorch/kernels/white_noise_kernel.py b/gpytorch/kernels/white_noise_kernel.py
index c9da7e47..149e4d28 100644
--- a/gpytorch/kernels/white_noise_kernel.py
+++ b/gpytorch/kernels/white_noise_kernel.py
@@ -50,7 +50,7 @@ class WhiteNoiseKernel(Kernel):
             self.is_batch = False
             self.variances = variances.unsqueeze(0).unsqueeze(-1)
 
-    def forward(self, x1, x2):
+    def forward(self, x1, x2, **params):
         if self.training and torch.equal(x1, x2):
             # Reshape into a batch of batch_size diagonal matrices, each of which is
             # (data_size * task_size) x (data_size * task_size)
diff --git a/gpytorch/lazy/interpolated_lazy_tensor.py b/gpytorch/lazy/interpolated_lazy_tensor.py
index d6afa8f7..c9ee0697 100644
--- a/gpytorch/lazy/interpolated_lazy_tensor.py
+++ b/gpytorch/lazy/interpolated_lazy_tensor.py
@@ -34,11 +34,23 @@ class InterpolatedLazyTensor(LazyTensor):
                 left_interp_indices = left_interp_indices.unsqueeze(0).expand(base_lazy_tensor.size(0), n_rows, 1)
             elif right_interp_indices is not None and right_interp_indices.ndimension() == 3:
                 left_interp_indices = left_interp_indices.unsqueeze(0).expand(right_interp_indices.size(0), n_rows, 1)
+        else:
+            if left_interp_indices.dim() != base_lazy_tensor.ndimension():
+                raise RuntimeError(
+                    "Expected left_interp_indices ({}) to have the same number of dimensions as "
+                    "base_lazy_Tensor ({})".format(left_interp_indices.size(), base_lazy_tensor.size())
+                )
 
         if left_interp_values is None:
             left_interp_values = torch.ones(
                 left_interp_indices.size(), dtype=base_lazy_tensor.dtype, device=base_lazy_tensor.device
             )
+        else:
+            if left_interp_indices.size() != left_interp_values.size():
+                raise RuntimeError(
+                    "Expected left_interp_indices ({}) to have the same size as "
+                    "left_interp_values ({})".format(left_interp_indices.size(), left_interp_values.size())
+                )
 
         if right_interp_indices is None:
             n_rows = base_lazy_tensor.size()[-2]
@@ -48,11 +60,23 @@ class InterpolatedLazyTensor(LazyTensor):
                 right_interp_indices = right_interp_indices.unsqueeze(0).expand(base_lazy_tensor.size(0), n_rows, 1)
             elif left_interp_indices.ndimension() == 3:
                 right_interp_indices = right_interp_indices.unsqueeze(0).expand(left_interp_indices.size(0), n_rows, 1)
+        else:
+            if left_interp_indices.dim() != base_lazy_tensor.ndimension():
+                raise RuntimeError(
+                    "Expected left_interp_indices ({}) to have the same number of dimensions as "
+                    "base_lazy_Tensor ({})".format(left_interp_indices.size(), base_lazy_tensor.size())
+                )
 
         if right_interp_values is None:
             right_interp_values = torch.ones(
                 right_interp_indices.size(), dtype=base_lazy_tensor.dtype, device=base_lazy_tensor.device
             )
+        else:
+            if left_interp_indices.size() != left_interp_values.size():
+                raise RuntimeError(
+                    "Expected left_interp_indices ({}) to have the same size as "
+                    "left_interp_values ({})".format(left_interp_indices.size(), left_interp_values.size())
+                )
 
         super(InterpolatedLazyTensor, self).__init__(
             base_lazy_tensor, left_interp_indices, left_interp_values, right_interp_indices, right_interp_values
@@ -324,7 +348,7 @@ class InterpolatedLazyTensor(LazyTensor):
             right_interp_indices = right_interp_indices.contiguous().view(-1)
 
             if self.base_lazy_tensor.ndimension() == 3:
-                batch_indices = torch.arange(0, batch_size, dtype=torch.long, device=self.device).unsqueeze_(-1)
+                batch_indices = torch.arange(0, batch_size, dtype=torch.long, device=self.device).unsqueeze(-1)
                 batch_indices = batch_indices.repeat(1, n_data * n_interp * n_interp).view(-1)
                 base_var_vals = self.base_lazy_tensor._batch_get_indices(
                     batch_indices, left_interp_indices, right_interp_indices
@@ -573,7 +597,7 @@ class InterpolatedLazyTensor(LazyTensor):
 
         # Increase interpolation indices appropriately
         factor = torch.arange(0, left_interp_indices.size(-3), dtype=torch.long, device=self.device)
-        factor = factor.unsqueeze_(-1).unsqueeze_(-1)
+        factor = factor.unsqueeze(-1).unsqueeze(-1)
         factor = factor * self.base_lazy_tensor.size(-1)
         if sum_batch_size is not None:
             factor = factor.unsqueeze(0)
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index 7c2f0a63..75c1bd5f 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -3,9 +3,11 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
+import torch
 from .lazy_tensor import LazyTensor
 from .non_lazy_tensor import NonLazyTensor
 from .lazy_tensor_representation_tree import LazyTensorRepresentationTree
+from .. import settings
 
 
 LAZY_KERNEL_TENSOR_WARNING = (
@@ -14,11 +16,12 @@ LAZY_KERNEL_TENSOR_WARNING = (
 
 
 class LazyEvaluatedKernelTensor(LazyTensor):
-    def __init__(self, kernel, x1, x2, squeeze_row=False, squeeze_col=False, **params):
+    def __init__(self, kernel, x1, x2, batch_dims=None, squeeze_row=False, squeeze_col=False, **params):
         super(LazyEvaluatedKernelTensor, self).__init__(kernel, x1, x2, **params)
         self.kernel = kernel
         self.x1 = x1
         self.x2 = x2
+        self.batch_dims = batch_dims
         self.squeeze_row = squeeze_row
         self.squeeze_col = squeeze_col
         self.is_batch = self.x1.ndimension() == 3 or (self.x1.ndimension() == 2 and self.squeeze_row)
@@ -62,6 +65,8 @@ class LazyEvaluatedKernelTensor(LazyTensor):
         Implementing it this way allows us to compute predictions more efficiently
         in cases where only the variances are required.
         """
+        from ..kernels import Kernel
+
         if hasattr(self, "_cached_kernel_diag"):
             return self._cached_kernel_diag
         elif hasattr(self, "_cached_kernel_eval"):
@@ -80,7 +85,40 @@ class LazyEvaluatedKernelTensor(LazyTensor):
             if x2.dim() == 2:  # We only have a single data point
                 x2 = x2.unsqueeze(1)
 
-            res = self.kernel.forward_diag(x1, x2, **self.params)
+            res = super(Kernel, self.kernel).__call__(x1, x2, diag=True, batch_dims=self.batch_dims, **self.params)
+
+            # Did this Kernel eat the diag option?
+            # If it does not return a LazyEvaluatedKernelTensor, we can call diag on the output
+            if not isinstance(res, LazyEvaluatedKernelTensor):
+                if res.dim() == x1.dim() and res.shape[-2:] == torch.Size((x1.size(-2), x2.size(-2))):
+                    res = res.diag()
+
+            # Now we'll make sure that the shape we're getting from diag makes sense
+            if settings.debug.on():
+                # If we used batch_dims...
+                shape = self.kernel.size(x1, x2)
+                if self.batch_dims == (0, 2):
+                    if len(shape) == 2:
+                        expected_shape = torch.Size((x1.size(-1), shape[0]))
+                    else:
+                        expected_shape = torch.Size((shape[0] * x1.size(-1), shape[1]))
+                    if res.shape != expected_shape:
+                        raise RuntimeError(
+                            "The kernel {} is not equipped to handle batch_dims=(0, 2) "
+                            "and diag. Expected size {}. Got size {}.".format(
+                                self.__class__.__name__, expected_shape, res.shape
+                            )
+                        )
+
+                # If we didn't use batch_dims...
+                else:
+                    expected_shape = shape[:-1]
+                    if res.shape != expected_shape:
+                        raise RuntimeError(
+                            "The kernel {} is not equipped to handle and diag. Expected size {}. "
+                            "Got size {}".format(self.__class__.__name__, expected_shape, res.shape)
+                        )
+
             if isinstance(res, LazyTensor):
                 res = res.evaluate()
             self._cached_kernel_diag = res.view(self.shape[:-1]).contiguous()
@@ -103,13 +141,16 @@ class LazyEvaluatedKernelTensor(LazyTensor):
                 x1 = self.x1
                 x2 = self.x2
 
-            self._cached_kernel_eval = super(Kernel, self.kernel).__call__(x1, x2, **self.params)
+            self._cached_kernel_eval = super(Kernel, self.kernel).__call__(
+                x1, x2, diag=False, batch_dims=self.batch_dims, **self.params
+            )
             if self.squeeze_row:
                 self._cached_kernel_eval.squeeze_(-2)
             if self.squeeze_col:
                 self._cached_kernel_eval.squeeze_(-1)
 
-            if not self.is_batch and self._cached_kernel_eval.ndimension() == 3:
+            if not self.is_batch and self._cached_kernel_eval.ndimension() == 3 \
+                    and self._cached_kernel_eval.size(0) == 1:
                 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)
@@ -198,4 +239,10 @@ class LazyEvaluatedKernelTensor(LazyTensor):
             )
 
     def _size(self):
-        return self.kernel.size(self.x1, self.x2)
+        size = self.kernel.size(self.x1, self.x2)
+        if self.batch_dims == (0, 2):
+            if len(size) == 2:
+                return torch.Size((self.x1.size(-1), size[0], size[1]))
+            else:
+                return torch.Size((self.x1.size(-1) * size[0], size[1], size[2]))
+        return size
diff --git a/gpytorch/lazy/sum_batch_lazy_tensor.py b/gpytorch/lazy/sum_batch_lazy_tensor.py
index 15a9155b..0f12178a 100644
--- a/gpytorch/lazy/sum_batch_lazy_tensor.py
+++ b/gpytorch/lazy/sum_batch_lazy_tensor.py
@@ -126,6 +126,12 @@ class SumBatchLazyTensor(BlockLazyTensor):
             res = res.sum(0)
         return res
 
+    def diag(self):
+        diag = self.base_lazy_tensor.diag()
+        if self.num_blocks is not None:
+            diag = diag.view(-1, self.num_blocks, diag.size(-1))
+        return diag.sum(-2)
+
     def zero_mean_mvn_samples(self, num_samples):
         num_dim = self.size(-2)
         res = self.base_lazy_tensor.zero_mean_mvn_samples(num_samples)
diff --git a/gpytorch/models/additive_grid_inducing_variational_gp.py b/gpytorch/models/additive_grid_inducing_variational_gp.py
index f282b1e8..52a1a82e 100644
--- a/gpytorch/models/additive_grid_inducing_variational_gp.py
+++ b/gpytorch/models/additive_grid_inducing_variational_gp.py
@@ -8,31 +8,31 @@ from ..priors import SmoothedBoxPrior
 
 
 class AdditiveGridInducingVariationalGP(GridInducingVariationalGP):
-    def __init__(self, grid_size, grid_bounds, n_components, mixing_params=False, sum_output=True):
+    def __init__(self, grid_size, grid_bounds, num_dim, mixing_params=False, sum_output=True):
         super(AdditiveGridInducingVariationalGP, self).__init__(grid_size, grid_bounds)
-        self.n_components = n_components
+        self.num_dim = num_dim
         self.sum_output = sum_output
 
         # Resize variational parameters to have one size per component
         variational_mean = self.variational_mean
         chol_variational_covar = self.chol_variational_covar
-        variational_mean.data.resize_(*([n_components] + list(variational_mean.size())))
-        chol_variational_covar.data.resize_(*([n_components] + list(chol_variational_covar.size())))
+        variational_mean.data.resize_(*([num_dim] + list(variational_mean.size())))
+        chol_variational_covar.data.resize_(*([num_dim] + list(chol_variational_covar.size())))
 
         # Mixing parameters
         if mixing_params:
             self.register_parameter(
                 name="mixing_params",
-                parameter=torch.nn.Parameter(torch.ones(n_components) / n_components),
-                prior=SmoothedBoxPrior(-2, 2, sigma=0.01, size=n_components),
+                parameter=torch.nn.Parameter(torch.ones(num_dim) / num_dim),
+                prior=SmoothedBoxPrior(-2, 2, sigma=0.01, size=num_dim),
             )
 
     def _compute_grid(self, inputs):
-        n_data, n_components, n_dimensions = inputs.size()
-        inputs = inputs.transpose(0, 1).contiguous().view(n_components * n_data, n_dimensions)
+        num_data, num_dim = inputs.size()
+        inputs = inputs.transpose(0, 1).contiguous().view(-1, 1)
         interp_indices, interp_values = super(AdditiveGridInducingVariationalGP, self)._compute_grid(inputs)
-        interp_indices = interp_indices.view(n_components, n_data, -1)
-        interp_values = interp_values.view(n_components, n_data, -1)
+        interp_indices = interp_indices.view(num_dim, num_data, -1)
+        interp_values = interp_values.view(num_dim, num_data, -1)
 
         if hasattr(self, "mixing_params"):
             interp_values = interp_values.mul(self.mixing_params.unsqueeze(1).unsqueeze(2))
@@ -51,33 +51,25 @@ class AdditiveGridInducingVariationalGP(GridInducingVariationalGP):
         self.chol_variational_covar.data.copy_(chol_covar_init)
 
     def covar_diag(self, inputs):
-        n_data, n_components, n_dimensions = inputs.size()
-        inputs = inputs.transpose(0, 1).contiguous()
+        num_data, num_dim = inputs.size()
+        inputs = inputs.transpose(0, 1).contiguous().unsqueeze(-1)
         return super(AdditiveGridInducingVariationalGP, self).covar_diag(inputs)
 
     def prior_output(self):
         out = super(AdditiveGridInducingVariationalGP, self).prior_output()
         mean = out.mean
-        covar = out.lazy_covariance_matrix.repeat(self.n_components, 1, 1)
+        covar = out.lazy_covariance_matrix.repeat(self.num_dim, 1, 1)
         return MultivariateNormal(mean, covar)
 
     def __call__(self, inputs, **kwargs):
         if inputs.ndimension() == 1:
-            inputs = inputs.unsqueeze(-1).unsqueeze(-1)
-        elif inputs.ndimension() == 2:
             inputs = inputs.unsqueeze(-1)
-        elif inputs.ndimension() != 3:
-            raise RuntimeError("AdditiveGridInducingVariationalGP expects a 3d tensor.")
-
-        n_data, n_components, n_dimensions = inputs.size()
-        if n_dimensions != self.grid.size(0):
-            raise RuntimeError("The number of dimensions should match the inducing points " " number of dimensions.")
-        if n_components != self.n_components:
-            raise RuntimeError("The number of components should match the number specified.")
-        if n_dimensions != 1:
-            raise RuntimeError(
-                "At the moment, AdditiveGridInducingVariationalGP only supports " "1d (Toeplitz) interpolation."
-            )
+        elif inputs.ndimension() != 2:
+            raise RuntimeError("AdditiveGridInducingVariationalGP expects a 2d tensor.")
+
+        num_data, num_dim = inputs.size()
+        if num_dim != self.num_dim:
+            raise RuntimeError("The number of dims should match the number specified.")
 
         output = super(AdditiveGridInducingVariationalGP, self).__call__(inputs, **kwargs)
         if self.sum_output:
diff --git a/gpytorch/models/grid_inducing_variational_gp.py b/gpytorch/models/grid_inducing_variational_gp.py
index b87d5020..25f240bc 100644
--- a/gpytorch/models/grid_inducing_variational_gp.py
+++ b/gpytorch/models/grid_inducing_variational_gp.py
@@ -19,16 +19,16 @@ class GridInducingVariationalGP(AbstractVariationalGP):
         self._grid_mode = True
         self._kernels = set()
 
-        grid = torch.zeros(len(grid_bounds), grid_size)
+        grid = torch.zeros(grid_size, len(grid_bounds))
         for i in range(len(grid_bounds)):
             grid_diff = float(grid_bounds[i][1] - grid_bounds[i][0]) / (grid_size - 2)
-            grid[i] = torch.linspace(grid_bounds[i][0] - grid_diff, grid_bounds[i][1] + grid_diff, grid_size)
+            grid[:, i] = torch.linspace(grid_bounds[i][0] - grid_diff, grid_bounds[i][1] + grid_diff, grid_size)
 
         inducing_points = torch.zeros(int(pow(grid_size, len(grid_bounds))), len(grid_bounds))
         prev_points = None
         for i in range(len(grid_bounds)):
             for j in range(grid_size):
-                inducing_points[j * grid_size ** i : (j + 1) * grid_size ** i, i].fill_(grid[i, j])
+                inducing_points[j * grid_size ** i : (j + 1) * grid_size ** i, i].fill_(grid[j, i])
                 if prev_points is not None:
                     inducing_points[j * grid_size ** i : (j + 1) * grid_size ** i, :i].copy_(prev_points)
             prev_points = inducing_points[: grid_size ** (i + 1), : (i + 1)]
@@ -104,7 +104,7 @@ class GridInducingVariationalGP(AbstractVariationalGP):
     def __getattr__(self, name):
         res = super(GridInducingVariationalGP, self).__getattr__(name)
         if name in self._kernels and not self._grid_mode:
-            return res.base_kernel_module
+            return res.base_kernel
         return res
 
     def __setattr__(self, name, value):
diff --git a/gpytorch/utils/interpolation.py b/gpytorch/utils/interpolation.py
index 9ba1e192..d9e1050a 100644
--- a/gpytorch/utils/interpolation.py
+++ b/gpytorch/utils/interpolation.py
@@ -42,8 +42,8 @@ class Interpolation(object):
 
     def interpolate(self, x_grid, x_target, interp_points=range(-2, 2)):
         # Do some boundary checking
-        grid_mins = x_grid.min(1)[0]
-        grid_maxs = x_grid.max(1)[0]
+        grid_mins = x_grid.min(0)[0]
+        grid_maxs = x_grid.max(0)[0]
         x_target_min = x_target.min(0)[0]
         x_target_max = x_target.min(0)[0]
         lt_min_mask = (x_target_min - grid_mins).lt(-1e-7)
@@ -81,7 +81,7 @@ class Interpolation(object):
         interp_points = torch.tensor(interp_points, dtype=x_grid.dtype, device=x_grid.device)
         interp_points_flip = interp_points.flip(0)
 
-        num_grid_points = x_grid.size(1)
+        num_grid_points = x_grid.size(0)
         num_target_points = x_target.size(0)
         num_dim = x_target.size(-1)
         num_coefficients = len(interp_points)
@@ -94,9 +94,9 @@ class Interpolation(object):
         )
 
         for i in range(num_dim):
-            grid_delta = x_grid[i, 1] - x_grid[i, 0]
-            lower_grid_pt_idxs = torch.floor((x_target[:, i] - x_grid[i, 0]) / grid_delta).squeeze()
-            lower_pt_rel_dists = (x_target[:, i] - x_grid[i, 0]) / grid_delta - lower_grid_pt_idxs
+            grid_delta = x_grid[1, i] - x_grid[0, i]
+            lower_grid_pt_idxs = torch.floor((x_target[:, i] - x_grid[0, i]) / grid_delta).squeeze()
+            lower_pt_rel_dists = (x_target[:, i] - x_grid[0, i]) / grid_delta - lower_grid_pt_idxs
             lower_grid_pt_idxs = lower_grid_pt_idxs - interp_points.max()
             lower_grid_pt_idxs.detach_()
 
@@ -110,7 +110,7 @@ class Interpolation(object):
 
             if num_left > 0:
                 left_boundary_pts.squeeze_(1)
-                x_grid_first = x_grid[i, :num_coefficients].unsqueeze(1).t().expand(num_left, num_coefficients)
+                x_grid_first = x_grid[:num_coefficients, i].unsqueeze(1).t().expand(num_left, num_coefficients)
 
                 grid_targets = x_target.select(1, i)[left_boundary_pts].unsqueeze(1).expand(num_left, num_coefficients)
                 dists = torch.abs(x_grid_first - grid_targets)
@@ -126,7 +126,7 @@ class Interpolation(object):
 
             if num_right > 0:
                 right_boundary_pts.squeeze_(1)
-                x_grid_last = x_grid[i, -num_coefficients:].unsqueeze(1).t().expand(num_right, num_coefficients)
+                x_grid_last = x_grid[-num_coefficients:, i].unsqueeze(1).t().expand(num_right, num_coefficients)
 
                 grid_targets = x_target.select(1, i)[right_boundary_pts].unsqueeze(1)
                 grid_targets = grid_targets.expand(num_right, num_coefficients)


Test patch:
diff --git a/test/examples/test_kissgp_additive_classification.py b/test/examples/test_kissgp_additive_classification.py
index 45d03fb8..7d014f35 100644
--- a/test/examples/test_kissgp_additive_classification.py
+++ b/test/examples/test_kissgp_additive_classification.py
@@ -28,10 +28,12 @@ train_y = train_y.float() * 2 - 1
 
 class GPClassificationModel(gpytorch.models.AdditiveGridInducingVariationalGP):
     def __init__(self):
-        super(GPClassificationModel, self).__init__(grid_size=16, grid_bounds=[(-1, 1)], n_components=2)
+        super(GPClassificationModel, self).__init__(grid_size=16, grid_bounds=[(-1, 1)], num_dim=2)
         self.mean_module = ConstantMean(prior=SmoothedBoxPrior(-1e-5, 1e-5))
         self.covar_module = ScaleKernel(
-            RBFKernel(log_lengthscale_prior=SmoothedBoxPrior(exp(-5), exp(6), sigma=0.1, log_transform=True)),
+            RBFKernel(
+                ard_num_dims=2, log_lengthscale_prior=SmoothedBoxPrior(exp(-5), exp(6), sigma=0.1, log_transform=True))
+            ,
             log_outputscale_prior=SmoothedBoxPrior(exp(-5), exp(6), sigma=0.1, log_transform=True),
         )
 
diff --git a/test/examples/test_kissgp_additive_regression.py b/test/examples/test_kissgp_additive_regression.py
index 7f72c483..e1b098c3 100644
--- a/test/examples/test_kissgp_additive_regression.py
+++ b/test/examples/test_kissgp_additive_regression.py
@@ -11,7 +11,7 @@ import torch
 import unittest
 import gpytorch
 from torch import optim
-from gpytorch.kernels import RBFKernel, AdditiveGridInterpolationKernel, ScaleKernel
+from gpytorch.kernels import RBFKernel, AdditiveStructureKernel, GridInterpolationKernel, ScaleKernel
 from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.means import ZeroMean
 from gpytorch.priors import SmoothedBoxPrior
@@ -41,10 +41,13 @@ class GPRegressionModel(gpytorch.models.ExactGP):
         super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
         self.mean_module = ZeroMean()
         self.base_covar_module = ScaleKernel(
-            RBFKernel(log_lengthscale_prior=SmoothedBoxPrior(exp(-3), exp(3), sigma=0.1, log_transform=True))
+            RBFKernel(
+                ard_num_dims=2, log_lengthscale_prior=SmoothedBoxPrior(exp(-3), exp(3), sigma=0.1, log_transform=True)
+            )
         )
-        self.covar_module = AdditiveGridInterpolationKernel(
-            self.base_covar_module, grid_size=100, num_dims=2, dim_groups=2
+        self.covar_module = AdditiveStructureKernel(
+            GridInterpolationKernel(self.base_covar_module, grid_size=100, num_dims=2),
+            num_dims=2,
         )
 
     def forward(self, x):
@@ -87,12 +90,12 @@ class TestKISSGPAdditiveRegression(unittest.TestCase):
                     optimizer.n_iter += 1
                     optimizer.step()
 
-                for param in gp_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)
+                    for param in gp_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)
 
                 # Test the model
                 gp_model.eval()
diff --git a/test/examples/test_kissgp_kronecker_product_classification.py b/test/examples/test_kissgp_kronecker_product_classification.py
index 2fc18171..15504f95 100644
--- a/test/examples/test_kissgp_kronecker_product_classification.py
+++ b/test/examples/test_kissgp_kronecker_product_classification.py
@@ -34,7 +34,9 @@ class GPClassificationModel(gpytorch.models.GridInducingVariationalGP):
         super(GPClassificationModel, self).__init__(grid_size=8, grid_bounds=[(0, 3), (0, 3)])
         self.mean_module = ConstantMean(prior=SmoothedBoxPrior(-1e-5, 1e-5))
         self.covar_module = ScaleKernel(
-            RBFKernel(log_lengthscale_prior=SmoothedBoxPrior(exp(-2.5), exp(3), sigma=0.1, log_transform=True))
+            RBFKernel(
+                ard_num_dims=2, log_lengthscale_prior=SmoothedBoxPrior(exp(-2.5), exp(3), sigma=0.1, log_transform=True)
+            )
         )
 
     def forward(self, x):
diff --git a/test/examples/test_kissgp_kronecker_product_regression.py b/test/examples/test_kissgp_kronecker_product_regression.py
index d7140fc9..3d3244bc 100644
--- a/test/examples/test_kissgp_kronecker_product_regression.py
+++ b/test/examples/test_kissgp_kronecker_product_regression.py
@@ -45,7 +45,9 @@ class GPRegressionModel(gpytorch.models.ExactGP):
         super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
         self.mean_module = ConstantMean(prior=SmoothedBoxPrior(-1, 1))
         self.base_covar_module = ScaleKernel(
-            RBFKernel(log_lengthscale_prior=SmoothedBoxPrior(exp(-3), exp(3), sigma=0.1, log_transform=True))
+            RBFKernel(
+                ard_num_dims=2, log_lengthscale_prior=SmoothedBoxPrior(exp(-3), exp(3), sigma=0.1, log_transform=True)
+            )
         )
         self.covar_module = GridInterpolationKernel(self.base_covar_module, grid_size=64, num_dims=2)
 
diff --git a/test/examples/test_kissgp_multiplicative_regression.py b/test/examples/test_kissgp_multiplicative_regression.py
index c29e8a82..41e1bc1d 100644
--- a/test/examples/test_kissgp_multiplicative_regression.py
+++ b/test/examples/test_kissgp_multiplicative_regression.py
@@ -11,7 +11,7 @@ import torch
 import unittest
 import gpytorch
 from torch import optim
-from gpytorch.kernels import RBFKernel, MultiplicativeGridInterpolationKernel, ScaleKernel
+from gpytorch.kernels import RBFKernel, ProductStructureKernel, GridInterpolationKernel, ScaleKernel
 from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.means import ConstantMean
 from gpytorch.priors import SmoothedBoxPrior
@@ -46,8 +46,9 @@ class GPRegressionModel(gpytorch.models.ExactGP):
         self.base_covar_module = ScaleKernel(
             RBFKernel(log_lengthscale_prior=SmoothedBoxPrior(exp(-3), exp(3), sigma=0.1, log_transform=True))
         )
-        self.covar_module = MultiplicativeGridInterpolationKernel(
-            self.base_covar_module, grid_size=100, num_dims=2, dim_groups=2
+        self.covar_module = ProductStructureKernel(
+            GridInterpolationKernel(self.base_covar_module, grid_size=100, num_dims=2),
+            num_dims=2
         )
 
     def forward(self, x):
diff --git a/test/kernels/test_cosine_kernel.py b/test/kernels/test_cosine_kernel.py
index 426b17d0..34e2f185 100644
--- a/test/kernels/test_cosine_kernel.py
+++ b/test/kernels/test_cosine_kernel.py
@@ -11,52 +11,91 @@ from gpytorch.kernels import CosineKernel
 
 class TestCosineKernel(unittest.TestCase):
     def test_computes_periodic_function(self):
-        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
-        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
+        a = torch.tensor([[4, 1], [2, 2], [8, 0]], dtype=torch.float)
+        b = torch.tensor([[0, 0], [2, 1], [1, 0]], dtype=torch.float)
         period = 1
         kernel = CosineKernel().initialize(log_period_length=math.log(period))
         kernel.eval()
 
-        actual = torch.zeros(3, 2)
+        actual = torch.zeros(3, 3)
         for i in range(3):
-            for j in range(2):
-                actual[i, j] = torch.cos(math.pi * (a[i] - b[j]) / period)
+            for j in range(3):
+                actual[i, j] = torch.cos(math.pi * ((a[i] - b[j]) / period).norm(2, dim=-1))
 
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = torch.zeros(2, 3, 3)
+        for i in range(3):
+            for j in range(3):
+                for l in range(2):
+                    actual[l, i, j] = torch.cos(math.pi * ((a[i, l] - b[j, l]) / period))
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
     def test_batch(self):
         a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
-        b = torch.tensor([[0, 2], [-1, 2]], dtype=torch.float).view(2, 2, 1)
+        b = torch.tensor([[0, 2, 1], [-1, 2, 0]], dtype=torch.float).view(2, 3, 1)
         period = torch.tensor(1, dtype=torch.float).view(1, 1, 1)
         kernel = CosineKernel().initialize(log_period_length=torch.log(period))
         kernel.eval()
 
-        actual = torch.zeros(2, 3, 2)
+        actual = torch.zeros(2, 3, 3)
         for k in range(2):
             for i in range(3):
-                for j in range(2):
-                    actual[k, i, j] = torch.cos(math.pi * (a[k, i] - b[k, j]) / period)
+                for j in range(3):
+                    actual[k, i, j] = torch.cos(math.pi * ((a[k, i] - b[k, j]) / period))
 
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_batch_separate(self):
-        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
-        b = torch.tensor([[0, 2], [-1, 2]], dtype=torch.float).view(2, 2, 1)
+        a = torch.tensor([[[4, 1], [2, 2], [8, 0]], [[2, 5], [6, 1], [0, 1]]], dtype=torch.float)
+        b = torch.tensor([[[0, 0], [2, 1], [1, 0]], [[1, 1], [2, 3], [1, 0]]], dtype=torch.float)
         period = torch.tensor([1, 2], dtype=torch.float).view(2, 1, 1)
         kernel = CosineKernel(batch_size=2).initialize(log_period_length=torch.log(period))
         kernel.eval()
 
-        actual = torch.zeros(2, 3, 2)
+        actual = torch.zeros(2, 3, 3)
         for k in range(2):
             for i in range(3):
-                for j in range(2):
-                    actual[k, i, j] = torch.cos(math.pi * (a[k, i] - b[k, j]) / period[k])
+                for j in range(3):
+                    actual[k, i, j] = torch.cos(math.pi * ((a[k, i] - b[k, j]) / period[k]).norm(2, dim=-1))
 
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = torch.zeros(4, 3, 3)
+        for k in range(2):
+            for i in range(3):
+                for j in range(3):
+                    for l in range(2):
+                        actual[k * 2 + l, i, j] = torch.cos(math.pi * ((a[k, i, l] - b[k, j, l]) / period[k]))
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/kernels/test_linear_kernel.py b/test/kernels/test_linear_kernel.py
index ad9ab318..9ac5ae67 100644
--- a/test/kernels/test_linear_kernel.py
+++ b/test/kernels/test_linear_kernel.py
@@ -11,7 +11,7 @@ from gpytorch.kernels import LinearKernel
 class TestLinearKernel(unittest.TestCase):
     def test_computes_linear_function_rectangular(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
-        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
+        b = torch.tensor([0, 2, 1], dtype=torch.float).view(3, 1)
 
         kernel = LinearKernel(num_dimensions=1).initialize(offset=0, variance=1.0)
         kernel.eval()
@@ -19,15 +19,65 @@ class TestLinearKernel(unittest.TestCase):
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
     def test_computes_linear_function_square(self):
-        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
+        a = torch.tensor([[4, 1], [2, 0], [8, 3]], dtype=torch.float)
 
-        kernel = LinearKernel(num_dimensions=1).initialize(offset=0, variance=1.0)
+        offset = torch.randn(1, 1, 2)
+        kernel = LinearKernel(num_dimensions=2).initialize(offset=offset, variance=1.0)
         kernel.eval()
-        actual = 1 + torch.matmul(a, a.t())
+        actual = 1 + torch.matmul(a - offset[0], (a - offset[0]).t())
         res = kernel(a, a).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+        # diag
+        res = kernel(a, a).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        dim_group_a = a - offset[0]
+        dim_group_a = dim_group_a.permute(1, 0).contiguous().view(-1, 3)
+        actual = 1 + torch.mul(dim_group_a.unsqueeze(-1), dim_group_a.unsqueeze(-2))
+        res = kernel(a, a, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, a, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_computes_linear_function_square_batch(self):
+        a = torch.tensor([[[4, 1], [2, 0], [8, 3]], [[1, 1], [2, 1], [1, 3]]], dtype=torch.float)
+
+        offset = torch.randn(1, 1, 2)
+        kernel = LinearKernel(num_dimensions=2).initialize(offset=offset, variance=1.0)
+        kernel.eval()
+        actual = 1 + torch.matmul(a - offset, (a - offset).transpose(-1, -2))
+        res = kernel(a, a).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, a).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        dim_group_a = a - offset
+        dim_group_a = dim_group_a.permute(0, 2, 1).contiguous().view(-1, 3)
+        actual = 1 + torch.mul(dim_group_a.unsqueeze(-1), dim_group_a.unsqueeze(-2))
+        res = kernel(a, a, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, a, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/kernels/test_matern_kernel.py b/test/kernels/test_matern_kernel.py
index 75253d6f..33777657 100644
--- a/test/kernels/test_matern_kernel.py
+++ b/test/kernels/test_matern_kernel.py
@@ -50,19 +50,36 @@ class TestMaternKernel(unittest.TestCase):
 
     def test_ard(self):
         a = torch.tensor([[1, 2], [3, 4]], dtype=torch.float)
-        b = torch.tensor([1, 4], dtype=torch.float).view(1, 2)
+        b = torch.tensor([[1, 4], [1, 4]], dtype=torch.float)
         lengthscales = torch.tensor([1, 2], dtype=torch.float).view(1, 1, 2)
 
         kernel = MaternKernel(nu=2.5, ard_num_dims=2)
         kernel.initialize(log_lengthscale=torch.log(lengthscales))
         kernel.eval()
 
-        dist = torch.tensor([[1], [2]], dtype=torch.float)
+        dist = torch.tensor([[1, 1], [2, 2]], dtype=torch.float)
         dist.mul_(math.sqrt(5))
         actual = (dist ** 2 / 3 + dist + 1).mul(torch.exp(-dist))
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-3)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        dist = torch.tensor([[[0, 0], [2, 2]], [[1, 1], [0, 0]]], dtype=torch.float)
+        dist.mul_(math.sqrt(5))
+        actual = (dist ** 2 / 3 + dist + 1).mul(torch.exp(-dist))
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
     def test_ard_batch(self):
         a = torch.tensor([[[1, 2, 3], [2, 4, 3]], [[2, -1, 2], [2, -1, 0]]], dtype=torch.float)
         b = torch.tensor([[[1, 4, 3]], [[2, -1, 0]]], dtype=torch.float)
@@ -79,18 +96,45 @@ class TestMaternKernel(unittest.TestCase):
 
     def test_ard_separate_batch(self):
         a = torch.tensor([[[1, 2, 3], [2, 4, 3]], [[2, -1, 2], [2, -1, 0]]], dtype=torch.float)
-        b = torch.tensor([[[1, 4, 3]], [[2, -1, 0]]], dtype=torch.float)
+        b = torch.tensor([[[1, 4, 3]], [[2, -1, 0]]], dtype=torch.float).repeat(1, 2, 1)
         lengthscales = torch.tensor([[[1, 2, 1]], [[2, 1, 0.5]]], dtype=torch.float)
 
         kernel = MaternKernel(nu=2.5, batch_size=2, ard_num_dims=3)
         kernel.initialize(log_lengthscale=torch.log(lengthscales))
         kernel.eval()
 
-        dist = torch.tensor([[[1], [1]], [[4], [0]]], dtype=torch.float).mul_(math.sqrt(5))
+        dist = torch.tensor([[[1, 1], [1, 1]], [[4, 4], [0, 0]]], dtype=torch.float).mul_(math.sqrt(5))
         actual = (dist ** 2 / 3 + dist + 1).mul(torch.exp(-dist))
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-3)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        dist = torch.tensor(
+            [
+                [[0, 0], [1, 1]],
+                [[1, 1], [0, 0]],
+                [[0, 0], [0, 0]],
+                [[0, 0], [0, 0]],
+                [[0, 0], [0, 0]],
+                [[4, 4], [0, 0]],
+            ],
+            dtype=torch.float,
+        )
+        dist.mul_(math.sqrt(5))
+        actual = (dist ** 2 / 3 + dist + 1).mul(torch.exp(-dist))
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/kernels/test_rbf_kernel.py b/test/kernels/test_rbf_kernel.py
index 68660461..b7418bf4 100644
--- a/test/kernels/test_rbf_kernel.py
+++ b/test/kernels/test_rbf_kernel.py
@@ -12,82 +12,123 @@ from gpytorch.kernels import RBFKernel
 class TestRBFKernel(unittest.TestCase):
     def test_ard(self):
         a = torch.tensor([[[1, 2], [2, 4]]], dtype=torch.float)
-        b = torch.tensor([1, 3], dtype=torch.float).view(1, 1, 2)
+        b = torch.tensor([[[1, 3], [0, 4]]], dtype=torch.float)
         lengthscales = torch.tensor([1, 2], dtype=torch.float).view(1, 1, 2)
 
         kernel = RBFKernel(ard_num_dims=2)
         kernel.initialize(log_lengthscale=lengthscales.log())
         kernel.eval()
 
-        actual = (a - b).div_(lengthscales).pow(2).sum(dim=-1).mul_(-0.5).exp()
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        actual = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1).mul_(-0.5).exp()
         res = kernel(a, b).evaluate()
-        self.assertLess(torch.norm(res - actual.unsqueeze(-1)), 1e-5)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # Diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = scaled_a.transpose(-1, -2).unsqueeze(-1) - scaled_b.transpose(-1, -2).unsqueeze(-2)
+        actual = actual.pow(2).mul_(-0.5).exp().view(2, 2, 2)
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims and diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_ard_batch(self):
         a = torch.tensor([[[1, 2, 3], [2, 4, 0]], [[-1, 1, 2], [2, 1, 4]]], dtype=torch.float)
-        b = torch.tensor([[[1, 3, 1]], [[2, -1, 0]]], dtype=torch.float)
+        b = torch.tensor([[[1, 3, 1]], [[2, -1, 0]]], dtype=torch.float).repeat(1, 2, 1)
         lengthscales = torch.tensor([[[1, 2, 1]]], dtype=torch.float)
 
         kernel = RBFKernel(batch_size=2, ard_num_dims=3)
         kernel.initialize(log_lengthscale=lengthscales.log())
         kernel.eval()
 
-        actual = (a - b).div_(lengthscales).pow(2).sum(dim=-1).mul_(-0.5).exp()
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        actual = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1).mul_(-0.5).exp()
         res = kernel(a, b).evaluate()
-        self.assertLess(torch.norm(res - actual.unsqueeze(-1)), 1e-5)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = scaled_a.transpose(-1, -2).unsqueeze(-1) - scaled_b.transpose(-1, -2).unsqueeze(-2)
+        actual = actual.pow(2).mul_(-0.5).exp().view(6, 2, 2)
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims and diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_ard_separate_batch(self):
         a = torch.tensor([[[1, 2, 3], [2, 4, 0]], [[-1, 1, 2], [2, 1, 4]]], dtype=torch.float)
-        b = torch.tensor([[[1, 3, 1]], [[2, -1, 0]]], dtype=torch.float)
+        b = torch.tensor([[[1, 3, 1]], [[2, -1, 0]]], dtype=torch.float).repeat(1, 2, 1)
         lengthscales = torch.tensor([[[1, 2, 1]], [[2, 1, 0.5]]], dtype=torch.float)
 
         kernel = RBFKernel(batch_size=2, ard_num_dims=3)
         kernel.initialize(log_lengthscale=lengthscales.log())
         kernel.eval()
 
-        actual = (a - b).div_(lengthscales).pow(2).sum(dim=-1).mul_(-0.5).exp()
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        actual = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1).mul_(-0.5).exp()
         res = kernel(a, b).evaluate()
-        self.assertLess(torch.norm(res - actual.unsqueeze(-1)), 1e-5)
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_subset_active_compute_radial_basis_function(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
         a_p = torch.tensor([1, 2, 3], dtype=torch.float).view(3, 1)
         a = torch.cat((a, a_p), 1)
-        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
+        b = torch.tensor([0, 2, 4], dtype=torch.float).view(3, 1)
         lengthscale = 2
 
         kernel = RBFKernel(active_dims=[0])
         kernel.initialize(log_lengthscale=math.log(lengthscale))
         kernel.eval()
 
-        actual = torch.tensor([[16, 4], [4, 0], [64, 36]], dtype=torch.float)
+        actual = torch.tensor([[16, 4, 0], [4, 0, 4], [64, 36, 16]], dtype=torch.float)
         actual.mul_(-0.5).div_(lengthscale ** 2).exp_()
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
     def test_computes_radial_basis_function(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
-        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
+        b = torch.tensor([0, 2, 4], dtype=torch.float).view(3, 1)
         lengthscale = 2
 
         kernel = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
         kernel.eval()
 
-        actual = torch.tensor([[16, 4], [4, 0], [64, 36]], dtype=torch.float).mul_(-0.5).div_(lengthscale ** 2).exp()
+        actual = torch.tensor([[16, 4, 0], [4, 0, 4], [64, 36, 16]], dtype=torch.float)
+        actual.mul_(-0.5).div_(lengthscale ** 2).exp_()
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
-    def test_forward_diag(self):
-        a = torch.tensor([4, 2, 8], dtype=torch.float).view(1, 3, 1)
-        b = torch.tensor([2, 0, 6], dtype=torch.float).view(1, 3, 1)
-        lengthscale = 2
-
-        kernel = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
-        kernel.eval()
-
-        res = kernel.forward_diag(a, b).squeeze()
-        actual = torch.tensor([0.60653066, 0.60653066, 0.60653066])
-
+        # diag
+        res = kernel(a, b).diag()
+        actual = actual.diag()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_computes_radial_basis_function_gradient(self):
diff --git a/test/kernels/test_scale_kernel.py b/test/kernels/test_scale_kernel.py
index c65771bc..4e5eda53 100644
--- a/test/kernels/test_scale_kernel.py
+++ b/test/kernels/test_scale_kernel.py
@@ -3,45 +3,82 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import math
 import torch
 import unittest
 from gpytorch.kernels import RBFKernel, ScaleKernel
 
 
 class TestScaleKernel(unittest.TestCase):
-    def test_forward(self):
-        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
-        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
-        lengthscale = 2
+    def test_ard(self):
+        a = torch.tensor([[[1, 2], [2, 4]]], dtype=torch.float).repeat(2, 1, 1)
+        b = torch.tensor([[[1, 3], [0, 4]]], dtype=torch.float).repeat(2, 1, 1)
+        lengthscales = torch.tensor([1, 2], dtype=torch.float).view(1, 1, 2)
 
-        base_kernel = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
+        base_kernel = RBFKernel(ard_num_dims=2)
+        base_kernel.initialize(log_lengthscale=lengthscales.log())
         kernel = ScaleKernel(base_kernel)
         kernel.initialize(log_outputscale=torch.tensor([3], dtype=torch.float).log())
         kernel.eval()
 
-        actual = torch.tensor([[16, 4], [4, 0], [64, 36]], dtype=torch.float)
-        actual.mul_(-0.5).div_(lengthscale ** 2).exp_().mul_(3)
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        actual = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1).mul_(-0.5).exp()
+        actual.mul_(3)
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
-    def test_forward_batch_mode(self):
-        a = torch.tensor([4, 2, 8], dtype=torch.float).view(1, 3, 1).repeat(4, 1, 1)
-        b = torch.tensor([0, 2], dtype=torch.float).view(1, 2, 1).repeat(4, 1, 1)
-        lengthscale = 2
+        # Diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = scaled_a.transpose(-1, -2).unsqueeze(-1) - scaled_b.transpose(-1, -2).unsqueeze(-2)
+        actual = actual.pow(2).mul_(-0.5).exp().view(4, 2, 2)
+        actual.mul_(3)
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims and diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_ard_batch(self):
+        a = torch.tensor([[[1, 2, 3], [2, 4, 0]], [[-1, 1, 2], [2, 1, 4]]], dtype=torch.float)
+        b = torch.tensor([[[1, 3, 1]], [[2, -1, 0]]], dtype=torch.float).repeat(1, 2, 1)
+        lengthscales = torch.tensor([[[1, 2, 1]]], dtype=torch.float)
 
-        base_kernel = RBFKernel().initialize(log_lengthscale=math.log(lengthscale))
-        kernel = ScaleKernel(base_kernel, batch_size=4)
-        kernel.initialize(log_outputscale=torch.tensor([1, 2, 3, 4], dtype=torch.float).log())
+        base_kernel = RBFKernel(batch_size=2, ard_num_dims=3)
+        base_kernel.initialize(log_lengthscale=lengthscales.log())
+        kernel = ScaleKernel(base_kernel, batch_size=2)
+        kernel.initialize(log_outputscale=torch.tensor([1, 2], dtype=torch.float).log())
         kernel.eval()
 
-        base_actual = (
-            torch.tensor([[16, 4], [4, 0], [64, 36]], dtype=torch.float).mul_(-0.5).div_(lengthscale ** 2).exp()
-        )
-        actual = base_actual.unsqueeze(0).mul(torch.tensor([1, 2, 3, 4], dtype=torch.float).view(4, 1, 1))
+        scaled_a = a.div(lengthscales)
+        scaled_b = b.div(lengthscales)
+        actual = (scaled_a.unsqueeze(-2) - scaled_b.unsqueeze(-3)).pow(2).sum(dim=-1).mul_(-0.5).exp()
+        actual[1].mul_(2)
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = scaled_a.transpose(-1, -2).unsqueeze(-1) - scaled_b.transpose(-1, -2).unsqueeze(-2)
+        actual = actual.pow(2).mul_(-0.5).exp().view(6, 2, 2)
+        actual[3:].mul_(2)
+        res = kernel(a, b, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims and diag
+        res = kernel(a, b, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/kernels/test_spectral_mixture_kernel.py b/test/kernels/test_spectral_mixture_kernel.py
index 468de4ef..090b624b 100644
--- a/test/kernels/test_spectral_mixture_kernel.py
+++ b/test/kernels/test_spectral_mixture_kernel.py
@@ -11,34 +11,61 @@ from gpytorch.kernels import SpectralMixtureKernel
 
 class TestSpectralMixtureKernel(unittest.TestCase):
     def test_standard(self):
-        a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
-        b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
-        means = [1, 2]
-        scales = [0.5, 0.25]
+        a = torch.tensor([[0, 1], [2, 2], [2, 0]], dtype=torch.float)
+        means = torch.tensor([[1, 2], [2, 1]], dtype=torch.float)
+        scales = torch.tensor([[0.5, 0.25], [0.25, 0.5]], dtype=torch.float)
         weights = [4, 2]
-        kernel = SpectralMixtureKernel(num_mixtures=2)
+        kernel = SpectralMixtureKernel(num_mixtures=2, ard_num_dims=2)
         kernel.initialize(
             log_mixture_weights=torch.tensor([[4, 2]], dtype=torch.float).log(),
-            log_mixture_means=torch.tensor([[[[1]], [[2]]]], dtype=torch.float).log(),
-            log_mixture_scales=torch.tensor([[[[0.5]], [[0.25]]]], dtype=torch.float).log(),
+            log_mixture_means=torch.tensor([[[[1, 2]], [[2, 1]]]], dtype=torch.float).log(),
+            log_mixture_scales=torch.tensor([[[[0.5, 0.25]], [[0.25, 0.5]]]], dtype=torch.float).log(),
         )
         kernel.eval()
 
-        actual = torch.zeros(3, 2)
+        actual = torch.zeros(2, 3, 3, 2)
         for i in range(3):
-            for j in range(2):
+            for j in range(3):
                 for k in range(2):
-                    new_term = torch.cos(2 * math.pi * (a[i] - b[j]) * means[k])
-                    new_term *= torch.exp(-2 * (math.pi * (a[i] - b[j])) ** 2 * scales[k] ** 2)
-                    new_term *= weights[k]
-                    actual[i, j] += new_term.item()
+                    new_term = torch.cos(2 * math.pi * (a[i] - a[j]) * means[k])
+                    new_term *= torch.exp(-2 * (math.pi * (a[i] - a[j])) ** 2 * scales[k] ** 2)
+                    actual[k, i, j] = new_term
+        actual = actual.prod(-1)
+        actual[0].mul_(weights[0])
+        actual[1].mul_(weights[1])
+        actual = actual.sum(0)
 
-        res = kernel(a, b).evaluate()
+        res = kernel(a, a).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # diag
+        res = kernel(a, a).diag()
+        actual = actual.diag()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims
+        actual = torch.zeros(2, 3, 3, 2)
+        for i in range(3):
+            for j in range(3):
+                for k in range(2):
+                    new_term = torch.cos(2 * math.pi * (a[i] - a[j]) * means[k])
+                    new_term *= torch.exp(-2 * (math.pi * (a[i] - a[j])) ** 2 * scales[k] ** 2)
+                    actual[k, i, j] = new_term
+        actual[0].mul_(weights[0])
+        actual[1].mul_(weights[1])
+        actual = actual.sum(0)
+        actual = actual.permute(2, 0, 1)
+        res = kernel(a, a, batch_dims=(0, 2)).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+        # batch_dims + diag
+        res = kernel(a, a, batch_dims=(0, 2)).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
         self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_batch_separate(self):
         a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
-        b = torch.tensor([[0, 2], [-1, 2]], dtype=torch.float).view(2, 2, 1)
+        b = torch.tensor([[0, 2, 1], [-1, 2, 0]], dtype=torch.float).view(2, 3, 1)
         means = torch.tensor([[1, 2], [2, 3]], dtype=torch.float).view(2, 2, 1, 1)
         scales = torch.tensor([[0.5, 0.25], [0.25, 1]], dtype=torch.float).view(2, 2, 1, 1)
         weights = torch.tensor([[4, 2], [1, 2]], dtype=torch.float).view(2, 2)
@@ -48,11 +75,11 @@ class TestSpectralMixtureKernel(unittest.TestCase):
         )
         kernel.eval()
 
-        actual = torch.zeros(2, 3, 2)
+        actual = torch.zeros(2, 3, 3)
         for l in range(2):
             for k in range(2):
                 for i in range(3):
-                    for j in range(2):
+                    for j in range(3):
                         new_term = torch.cos(2 * math.pi * (a[l, i] - b[l, j]) * means[l, k])
                         new_term *= torch.exp(-2 * (math.pi * (a[l, i] - b[l, j])) ** 2 * scales[l, k] ** 2)
                         new_term *= weights[l, k]
@@ -61,6 +88,11 @@ class TestSpectralMixtureKernel(unittest.TestCase):
         res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
+        # diag
+        res = kernel(a, b).diag()
+        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/utils/test_interpolation.py b/test/utils/test_interpolation.py
index b46cc2ce..a1daf8d7 100644
--- a/test/utils/test_interpolation.py
+++ b/test/utils/test_interpolation.py
@@ -14,11 +14,11 @@ from gpytorch.utils.interpolation import Interpolation, left_interp, left_t_inte
 class TestCubicInterpolation(unittest.TestCase):
     def test_interpolation(self):
         x = torch.linspace(0.01, 1, 100).unsqueeze(1)
-        grid = torch.linspace(-0.05, 1.05, 50).unsqueeze(0)
+        grid = torch.linspace(-0.05, 1.05, 50).unsqueeze(1)
         indices, values = Interpolation().interpolate(grid, x)
         indices = indices.squeeze_(0)
         values = values.squeeze_(0)
-        test_func_grid = grid.squeeze(0).pow(2)
+        test_func_grid = grid.squeeze(1).pow(2)
         test_func_x = x.pow(2).squeeze(-1)
 
         interp_func_x = gpytorch.utils.interpolation.left_interp(indices, values, test_func_grid.unsqueeze(1)).squeeze()
@@ -27,7 +27,7 @@ class TestCubicInterpolation(unittest.TestCase):
 
     def test_multidim_interpolation(self):
         x = torch.tensor([[0.25, 0.45, 0.65, 0.85], [0.35, 0.375, 0.4, 0.425], [0.45, 0.5, 0.55, 0.6]]).t().contiguous()
-        grid = torch.linspace(0., 1., 11).unsqueeze(0).repeat(3, 1)
+        grid = torch.linspace(0., 1., 11).unsqueeze(1).repeat(1, 3)
 
         indices, values = Interpolation().interpolate(grid, x)