Repository: cornellius-gp/gpytorch
Pull request #605 (2019-03-31 15:58:48)

Remove default batch dimension from kernels
Right now the default batch shape for kernels is `torch.Size([1])`, with some squeezing if the data is not supplied as batch data.

This makes the default `batch_shape=torch.Size([])`. Saved models with the old default batch shape still load (with deprecation warning).

(Note that this speeds up evaluation quite a bit with small amounts of data).

Balandat: Thanks, this looks like the right thing to do.

> Note that this speeds up evaluation quite a bit with small amounts of data

Nice.

Patch:
diff --git a/gpytorch/distributions/multivariate_normal.py b/gpytorch/distributions/multivariate_normal.py
index ba15fb41..2942ebb3 100644
--- a/gpytorch/distributions/multivariate_normal.py
+++ b/gpytorch/distributions/multivariate_normal.py
@@ -116,10 +116,14 @@ class _MultivariateNormalBase(TMultivariateNormal, Distribution):
 
         # Repeat the covar to match the batch shape of diff
         if diff.shape[:-1] != covar.batch_shape:
-            padded_batch_shape = (*(1 for _ in range(diff.dim() + 1 - covar.dim())), *covar.batch_shape)
-            covar = covar.repeat(
-                *(diff_size // covar_size for diff_size, covar_size in zip(diff.shape[:-1], padded_batch_shape)), 1, 1
-            )
+            if len(diff.shape[:-1]) < len(covar.batch_shape):
+                diff = diff.expand(covar.shape[:-1])
+            else:
+                padded_batch_shape = (*(1 for _ in range(diff.dim() + 1 - covar.dim())), *covar.batch_shape)
+                covar = covar.repeat(
+                    *(diff_size // covar_size for diff_size, covar_size in zip(diff.shape[:-1], padded_batch_shape)),
+                    1, 1
+                )
 
         # Get log determininat and first part of quadratic form
         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
diff --git a/gpytorch/kernels/additive_structure_kernel.py b/gpytorch/kernels/additive_structure_kernel.py
index 4d9cec46..e79775a2 100644
--- a/gpytorch/kernels/additive_structure_kernel.py
+++ b/gpytorch/kernels/additive_structure_kernel.py
@@ -49,7 +49,7 @@ class AdditiveStructureKernel(Kernel):
             evaluate = True
             res = NonLazyTensor(res)
 
-        res = res.sum(-3).unsqueeze(0)
+        res = res.sum(0)
 
         if evaluate:
             res = res.evaluate()
diff --git a/gpytorch/kernels/cosine_kernel.py b/gpytorch/kernels/cosine_kernel.py
index cfb718b6..549692dc 100644
--- a/gpytorch/kernels/cosine_kernel.py
+++ b/gpytorch/kernels/cosine_kernel.py
@@ -23,7 +23,7 @@ class CosineKernel(Kernel):
     Args:
         :attr:`batch_shape` (torch.Size, optional):
             Set this if you want a separate lengthscale for each
-            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`
+            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([])`
         :attr:`active_dims` (tuple of ints, optional):
             Set this if you want to compute the covariance of only a few input dimensions. The ints
             corresponds to the indices of the dimensions. Default: `None`.
diff --git a/gpytorch/kernels/grid_interpolation_kernel.py b/gpytorch/kernels/grid_interpolation_kernel.py
index a56a9fdc..147e6a8d 100644
--- a/gpytorch/kernels/grid_interpolation_kernel.py
+++ b/gpytorch/kernels/grid_interpolation_kernel.py
@@ -114,23 +114,22 @@ class GridInterpolationKernel(GridKernel):
         )
 
     def _compute_grid(self, inputs, batch_dims):
-        batch_size, n_data, n_dimensions = inputs.size()
+        n_data, n_dimensions = inputs.size(-2), inputs.size(-1)
         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.unsqueeze(0).transpose(0, -1)
+            n_dimensions = 1
+        batch_shape = inputs.shape[:-2]
 
-        inputs = inputs.view(batch_size * n_data, n_dimensions)
+        inputs = inputs.contiguous().view(-1, 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)
+        interp_indices = interp_indices.view(*batch_shape, n_data, -1)
+        interp_values = interp_values.view(*batch_shape, n_data, -1)
         return interp_indices, interp_values
 
     def _inducing_forward(self, batch_dims, **params):
         return super(GridInterpolationKernel, self).forward(self.grid, self.grid, batch_dims=batch_dims, **params)
 
-    def forward(self, x1, x2, batch_dims=None, **params):
+    def forward(self, x1, x2, diag=False, 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):
@@ -160,9 +159,9 @@ class GridInterpolationKernel(GridKernel):
 
         base_lazy_tsr = self._inducing_forward(batch_dims=batch_dims, **params)
         if batch_dims == (0, 2):
-            base_lazy_tsr = base_lazy_tsr.repeat(x1.size(-1), 1, 1)
-        if x1.size(0) > 1:
-            base_lazy_tsr = base_lazy_tsr.repeat(x1.size(0), 1, 1)
+            base_lazy_tsr = base_lazy_tsr.repeat(x1.size(-1), *x1.shape[:-2], 1, 1)
+        if x1.dim() > 2:
+            base_lazy_tsr = base_lazy_tsr.repeat(*x1.shape[:-2], 1, 1)
 
         left_interp_indices, left_interp_values = self._compute_grid(x1, batch_dims)
         if torch.equal(x1, x2):
@@ -179,7 +178,10 @@ class GridInterpolationKernel(GridKernel):
             right_interp_values,
         )
 
-        return res
+        if diag:
+            return res.diag()
+        else:
+            return res
 
     def size(self, x1, x2):
         return self.base_kernel.size(x1, x2)
diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index 99615dfe..4e9ce80a 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -58,13 +58,13 @@ class GridKernel(Kernel):
         return self
 
     def forward(self, x1, x2, diag=False, batch_dims=None, **params):
-        grid = self.grid.unsqueeze(0)
+        grid = self.grid
 
         if not self.interpolation_mode:
-            full_grid = self.full_grid.unsqueeze(0) if self.full_grid.dim() == 2 else self.full_grid
-
-        x1 = x1.unsqueeze(0) if x1.dim() == 2 else x1
-        x2 = x2.unsqueeze(0) if x2.dim() == 2 else x2
+            if len(x1.shape[:-2]):
+                full_grid = self.full_grid.expand(*x1.shape[:-2], *self.full_grid.shape[-2:])
+            else:
+                full_grid = self.full_grid
 
         if self.interpolation_mode or (torch.equal(x1, full_grid) and torch.equal(x2, full_grid)):
             if not self.training and hasattr(self, "_cached_kernel_mat"):
@@ -73,19 +73,19 @@ class GridKernel(Kernel):
             n_dim = x1.size(-1)
 
             if settings.use_toeplitz.on():
-                first_item = grid[:, 0:1]
+                first_item = grid[0:1]
                 covar_columns = self.base_kernel(first_item, grid, diag=False, batch_dims=(0, 2), **params)
                 covar_columns = delazify(covar_columns).squeeze(-2)
                 if batch_dims == (0, 2):
-                    covars = [ToeplitzLazyTensor(covar_columns)]
+                    covars = [ToeplitzLazyTensor(covar_columns.squeeze(-2))]
                 else:
-                    covars = [ToeplitzLazyTensor(covar_columns[i : i + 1]) for i in range(n_dim)]
+                    covars = [ToeplitzLazyTensor(covar_columns[i : i + 1].squeeze(-2)) for i in range(n_dim)]
             else:
                 full_covar = self.base_kernel(grid, grid, batch_dims=(0, 2), **params)
                 if batch_dims == (0, 2):
                     covars = [full_covar]
                 else:
-                    covars = [full_covar[i : i + 1] for i in range(n_dim)]
+                    covars = [full_covar[i : i + 1].squeeze(0) for i in range(n_dim)]
 
             if len(covars) > 1:
                 covar = KroneckerProductLazyTensor(*covars[::-1])
diff --git a/gpytorch/kernels/inducing_point_kernel.py b/gpytorch/kernels/inducing_point_kernel.py
index a1d0ae7a..bb1ea952 100644
--- a/gpytorch/kernels/inducing_point_kernel.py
+++ b/gpytorch/kernels/inducing_point_kernel.py
@@ -3,10 +3,10 @@
 import math
 import torch
 from .kernel import Kernel
-from ..functions import add_jitter
 from ..lazy import delazify, DiagLazyTensor, MatmulLazyTensor, RootLazyTensor, PsdSumLazyTensor
 from ..distributions import MultivariateNormal
 from ..mlls import InducingPointKernelAddedLossTerm
+from ..utils.cholesky import psd_safe_cholesky
 
 
 class InducingPointKernel(Kernel):
@@ -19,7 +19,7 @@ class InducingPointKernel(Kernel):
             inducing_points = inducing_points.unsqueeze(-1)
         if inducing_points.ndimension() != 2:
             raise RuntimeError("Inducing points should be 2 dimensional")
-        self.register_parameter(name="inducing_points", parameter=torch.nn.Parameter(inducing_points.unsqueeze(0)))
+        self.register_parameter(name="inducing_points", parameter=torch.nn.Parameter(inducing_points))
         self.register_added_loss_term("inducing_point_loss_term")
 
     def train(self, mode=True):
@@ -42,11 +42,11 @@ class InducingPointKernel(Kernel):
         if not self.training and hasattr(self, "_cached_kernel_inv_root"):
             return self._cached_kernel_inv_root
         else:
-            chols = torch.cholesky(add_jitter(self._inducing_mat), upper=True)
-            eye = torch.eye(chols[0].size(-1), device=chols[0].device)
-            inv_root_list = [torch.trtrs(eye, chol)[0].unsqueeze(0) for chol in chols]
+            chol = psd_safe_cholesky(self._inducing_mat, upper=True)
+            eye = torch.eye(chol.size(-1), device=chol.device)
+            inv_root = torch.trtrs(eye, chol)[0]
 
-            res = torch.cat(inv_root_list, 0)
+            res = inv_root
             if not self.training:
                 self._cached_kernel_inv_root = res
             return res
@@ -70,17 +70,12 @@ class InducingPointKernel(Kernel):
     def _covar_diag(self, inputs):
         if inputs.ndimension() == 1:
             inputs = inputs.unsqueeze(1)
-        orig_size = list(inputs.size())
-
-        # Resize inputs so that everything is batch
-        inputs = inputs.unsqueeze(-2).view(-1, 1, inputs.size(-1))
 
         # Get diagonal of covar
-        covar_diag = delazify(self.base_kernel(inputs))
-        covar_diag = covar_diag.view(orig_size[:-1])
+        covar_diag = delazify(self.base_kernel(inputs, diag=True))
         return DiagLazyTensor(covar_diag)
 
-    def forward(self, x1, x2, **kwargs):
+    def forward(self, x1, x2, diag=False, **kwargs):
         covar = self._get_covariance(x1, x2)
 
         if self.training:
@@ -94,7 +89,10 @@ class InducingPointKernel(Kernel):
             )
             self.update_added_loss_term("inducing_point_loss_term", new_added_loss_term)
 
-        return covar
+        if diag:
+            return covar.diag()
+        else:
+            return covar
 
     def size(self, x1, x2):
         return self.base_kernel.size(x1, x2)
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 6c4fc72f..c7ab0d59 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -2,6 +2,7 @@
 
 from abc import abstractmethod
 import torch
+import warnings
 from torch.nn import ModuleList
 from ..lazy import lazify, delazify, LazyEvaluatedKernelTensor, ZeroLazyTensor
 from ..module import Module
@@ -9,7 +10,6 @@ from .. import settings
 from ..utils.deprecation import _deprecate_kwarg_with_transform
 from ..utils import broadcasting
 from ..constraints import Positive
-import warnings
 
 
 @torch.jit.script
@@ -101,7 +101,7 @@ class Kernel(Module):
             dimension. It should be `d` if :attr:`x1` is a `n x d` matrix.  Default: `None`
         :attr:`batch_shape` (torch.Size, optional):
             Set this if you want a separate lengthscale for each batch of input
-            data. It should be `b1 x ... x bk` if :attr:`x1` is a `b1 x ... x bk x n x d` tensor.  Default: `1`
+            data. It should be `b1 x ... x bk` if :attr:`x1` is a `b1 x ... x bk x n x d` tensor.
         :attr:`active_dims` (tuple of ints, optional):
             Set this if you want to compute the covariance of only a few input dimensions. The ints
             corresponds to the indices of the dimensions. Default: `None`.
@@ -124,11 +124,29 @@ class Kernel(Module):
         >>> tensor_covar_matrix = lazy_covar_matrix.evaluate() # Gets the actual tensor for this kernel matrix
     """
 
+    def _batch_shape_state_dict_hook(
+        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+    ):
+        if not len(self.batch_shape):
+            try:
+                current_state_dict = self.state_dict()
+                for name, param in current_state_dict.items():
+                    load_param = state_dict[prefix + name]
+                    if load_param.dim() == param.dim() + 1:
+                        warnings.warn(
+                            f"The supplied state_dict contains a parameter ({prefix + name}) with an extra batch "
+                            f"dimension ({load_param.shape} vs {param.shape}).\nDefault batch shapes are now "
+                            "deprecated in GPyTorch. You may wish to re-save your model.", DeprecationWarning
+                        )
+                        load_param.squeeze_(0)
+            except Exception:
+                pass
+
     def __init__(
         self,
         has_lengthscale=False,
         ard_num_dims=None,
-        batch_shape=torch.Size([1]),
+        batch_shape=torch.Size([]),
         active_dims=None,
         lengthscale_prior=None,
         lengthscale_constraint=None,
@@ -136,6 +154,8 @@ class Kernel(Module):
         **kwargs
     ):
         super(Kernel, self).__init__()
+        self._register_load_state_dict_pre_hook(self._batch_shape_state_dict_hook)
+
         if active_dims is not None and not torch.is_tensor(active_dims):
             active_dims = torch.tensor(active_dims, dtype=torch.long)
         self.register_buffer("active_dims", active_dims)
@@ -214,6 +234,13 @@ class Kernel(Module):
             error_msg="x1 and x2 were not broadcastable to a proper kernel shape."
             "Got x1.shape = {} and x2.shape = {}".format(str(x1.shape), str(x2.shape))
         )
+        expected_size = broadcasting._mul_broadcast_shape(
+            expected_size[:-2], self.batch_shape,
+            error_msg=(
+                f"x1 and x2 were not broadcastable with kernel of batch_shape {self.batch_shape}."
+                f"Got x1.shape = {x1.shape} and x2.shape = {x2.shape}"
+            )
+        ) + expected_size[-2:]
         return expected_size
 
     @abstractmethod
@@ -304,11 +331,7 @@ class Kernel(Module):
         if diag:
             # Special case the diagonal because we can return all zeros most of the time.
             if x1_eq_x2:
-                # We are computing distances between points and themselves, which are all zero
-                if batch_dims == (0, 2):
-                    res = torch.zeros(x1.shape[0] * x1.shape[-1], x2.shape[-2], dtype=x1.dtype, device=x1.device)
-                else:
-                    res = torch.zeros(x1.shape[0], x2.shape[-2], dtype=x1.dtype, device=x1.device)
+                res = torch.zeros(*x1.shape[:-2], x1.shape[-2], dtype=x1.dtype, device=x1.device)
                 if postprocess:
                     res = dist_postprocess_func(res)
                 return res
@@ -318,18 +341,13 @@ class Kernel(Module):
                     res = res.pow(2)
             if postprocess:
                 res = dist_postprocess_func(res)
+            return res
 
         elif square_dist:
             res = self.distance_module._jit_sq_dist(x1, x2, torch.tensor(x1_eq_x2), postprocess)
         else:
             res = self.distance_module._jit_dist(x1, x2, torch.tensor(x1_eq_x2), postprocess, torch.tensor(False))
 
-        if batch_dims == (0, 2):
-            if diag:
-                res = res.transpose(0, 1).contiguous().view(-1, res.shape[-1])
-            else:
-                res = res.transpose(0, 1).contiguous().view(-1, *res.shape[-2:])
-
         return res
 
     def __call__(self, x1, x2=None, diag=False, batch_dims=None, **params):
@@ -350,13 +368,6 @@ class Kernel(Module):
             if not x1_.size(-1) == x2_.size(-1):
                 raise RuntimeError("x1_ and x2_ must have the same number of dimensions!")
 
-        # TODO: remove bach checking once kernels support arbitrary batch dimensions
-        is_batch = x1_.dim() > 2
-        if not is_batch:
-            x1_ = x1_.unsqueeze(0)
-            if x2_ is not None:
-                x2_ = x2_.unsqueeze(0)
-
         if x2_ is None:
             x2_ = x1_
 
@@ -378,13 +389,6 @@ class Kernel(Module):
         if diag:
             res = super(Kernel, self).__call__(x1_, x2_, diag=True, batch_dims=batch_dims, **params)
 
-            # TODO: remove bach checking once kernels support arbitrary batch dimensions
-            if not is_batch:
-                res = res.squeeze(0)
-                x1_ = x1_.squeeze(0)
-                if x2_ is not x1_:
-                    x2_ = x2_.squeeze(0)
-
             # 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):
@@ -396,10 +400,7 @@ class Kernel(Module):
                 # 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]))
+                    expected_shape = torch.Size((x1.size(-1),) + shape[:-1])
                     if res.shape != expected_shape:
                         raise RuntimeError(
                             "The kernel {} is not equipped to handle batch_dims=(0, 2) "
@@ -424,22 +425,12 @@ class Kernel(Module):
             else:
                 res = super(Kernel, self).__call__(x1_, x2_, batch_dims=batch_dims, **params)
 
-            # TODO: remove bach checking once kernels support arbitrary batch dimensions
-            if not is_batch:
-                res = res.squeeze(0)
-                x1_ = x1_.squeeze(0)
-                if x2_ is not x1_:
-                    x2_ = x2_.squeeze(0)
-
             # 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]))
+                    expected_shape = torch.Size((x1_.size(-1),) + shape)
                     if res.shape != expected_shape:
                         raise RuntimeError(
                             "The kernel {} is not equipped to handle batch_dims=(0, 2). Expected size {}. "
diff --git a/gpytorch/kernels/linear_kernel.py b/gpytorch/kernels/linear_kernel.py
index 80e2355c..21bc324d 100644
--- a/gpytorch/kernels/linear_kernel.py
+++ b/gpytorch/kernels/linear_kernel.py
@@ -97,9 +97,7 @@ class LinearKernel(Kernel):
     def forward(self, x1, x2, diag=False, batch_dims=None, **params):
         x1_ = x1 * self.variance.sqrt()
         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))
+            x1_ = x1_.unsqueeze(0).transpose(0, -1)
 
         if x1.size() == x2.size() and torch.equal(x1, x2):
             # Use RootLazyTensor when x1 == x2 for efficiency when composing
@@ -109,10 +107,11 @@ class LinearKernel(Kernel):
         else:
             x2_ = x2 * self.variance.sqrt()
             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))
+                x2_ = x2_.unsqueeze(0).transpose(0, -1)
 
-            prod = MatmulLazyTensor(x1_, x2_.transpose(2, 1))
+            prod = MatmulLazyTensor(x1_, x2_.transpose(-2, -1))
 
-        return prod
+        if diag:
+            return prod.diag()
+        else:
+            return prod
diff --git a/gpytorch/kernels/matern_kernel.py b/gpytorch/kernels/matern_kernel.py
index 644a138c..37952152 100644
--- a/gpytorch/kernels/matern_kernel.py
+++ b/gpytorch/kernels/matern_kernel.py
@@ -42,7 +42,7 @@ class MaternKernel(Kernel):
             input dimension. It should be `d` if :attr:`x1` is a `n x d` matrix. Default: `None`
         :attr:`batch_shape` (torch.Size, optional):
             Set this if you want a separate lengthscale for each
-             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`
+             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([])`
         :attr:`active_dims` (tuple of ints, optional):
             Set this if you want to
             compute the covariance of only a few input dimensions. The ints
@@ -81,18 +81,18 @@ class MaternKernel(Kernel):
         super(MaternKernel, self).__init__(has_lengthscale=True, **kwargs)
         self.nu = nu
 
-    def forward(self, x1, x2, **params):
+    def forward(self, x1, x2, diag=False, **params):
         if (
             x1.requires_grad
             or x2.requires_grad
             or (self.ard_num_dims is not None and self.ard_num_dims > 1)
-            or params.get('diag', False)
+            or diag
         ):
             mean = x1.contiguous().view(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
 
             x1_ = (x1 - mean).div(self.lengthscale)
             x2_ = (x2 - mean).div(self.lengthscale)
-            distance = self._covar_dist(x1_, x2_, **params)
+            distance = self._covar_dist(x1_, x2_, diag=diag, **params)
             exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)
 
             if self.nu == 0.5:
diff --git a/gpytorch/kernels/multitask_kernel.py b/gpytorch/kernels/multitask_kernel.py
index 1cabd2cd..73c61b18 100644
--- a/gpytorch/kernels/multitask_kernel.py
+++ b/gpytorch/kernels/multitask_kernel.py
@@ -43,7 +43,8 @@ class MultitaskKernel(Kernel):
         if batch_dims == (0, 2):
             raise RuntimeError("AdditiveGridInterpolationKernel does not accept the batch_dims argument.")
         covar_i = self.task_covar_module.covar_matrix
-        covar_i = covar_i.repeat(x1.size(0), 1, 1)
+        if len(x1.shape[:-2]):
+            covar_i = covar_i.repeat(*x1.shape[:-2], 1, 1)
         covar_x = lazify(self.data_covar_module.forward(x1, x2, **params))
         res = KroneckerProductLazyTensor(covar_x, covar_i)
         return res.diag() if diag else res
@@ -55,6 +56,6 @@ class MultitaskKernel(Kernel):
         """
         non_batch_size = (self.num_tasks * x1.size(-2), self.num_tasks * x2.size(-2))
         if x1.ndimension() == 3:
-            return torch.Size((x1.size(0),) + non_batch_size)
+            return torch.Size(x1.shape[:-2] + non_batch_size)
         else:
             return torch.Size(non_batch_size)
diff --git a/gpytorch/kernels/periodic_kernel.py b/gpytorch/kernels/periodic_kernel.py
index ec5dad50..d8566156 100644
--- a/gpytorch/kernels/periodic_kernel.py
+++ b/gpytorch/kernels/periodic_kernel.py
@@ -35,7 +35,7 @@ class PeriodicKernel(Kernel):
     Args:
         :attr:`batch_shape` (torch.Size, optional):
             Set this if you want a separate lengthscale for each
-             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.
+             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([])`.
         :attr:`active_dims` (tuple of ints, optional):
             Set this if you want to compute the covariance of only a few input dimensions. The ints
             corresponds to the indices of the dimensions. Default: `None`.
diff --git a/gpytorch/kernels/product_structure_kernel.py b/gpytorch/kernels/product_structure_kernel.py
index 42d3b163..0f24618a 100644
--- a/gpytorch/kernels/product_structure_kernel.py
+++ b/gpytorch/kernels/product_structure_kernel.py
@@ -54,7 +54,7 @@ class ProductStructureKernel(Kernel):
             evaluate = True
             res = NonLazyTensor(res)
 
-        res = res.prod(-3).unsqueeze(0)
+        res = res.prod(0)
 
         if evaluate:
             res = res.evaluate()
diff --git a/gpytorch/kernels/rbf_kernel.py b/gpytorch/kernels/rbf_kernel.py
index b6ddc84a..d63de39a 100644
--- a/gpytorch/kernels/rbf_kernel.py
+++ b/gpytorch/kernels/rbf_kernel.py
@@ -36,7 +36,7 @@ class RBFKernel(Kernel):
             input dimension. It should be `d` if :attr:`x1` is a `n x d` matrix. Default: `None`
         :attr:`batch_shape` (torch.Size, optional):
             Set this if you want a separate lengthscale for each
-            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.
+            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([])`.
         :attr:`active_dims` (tuple of ints, optional):
             Set this if you want to compute the covariance of only a few input dimensions. The ints
             corresponds to the indices of the dimensions. Default: `None`.
diff --git a/gpytorch/kernels/rbf_kernel_grad.py b/gpytorch/kernels/rbf_kernel_grad.py
index 6d3c6c85..4c60755a 100644
--- a/gpytorch/kernels/rbf_kernel_grad.py
+++ b/gpytorch/kernels/rbf_kernel_grad.py
@@ -20,7 +20,7 @@ class RBFKernelGrad(RBFKernel):
     Args:
         :attr:`batch_shape` (torch.Size, optional):
             Set this if you want a separate lengthscale for each
-             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.
+             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([])`.
         :attr:`active_dims` (tuple of ints, optional):
             Set this if you want to compute the covariance of only a few input dimensions. The ints
             corresponds to the indices of the dimensions. Default: `None`.
@@ -50,15 +50,11 @@ class RBFKernelGrad(RBFKernel):
         >>> covar = covar_module(x)  # Output: LazyTensor of size (2 x 60 x 60)
     """
     def forward(self, x1, x2, diag=False, **params):
-        b = 1
-        if len(x1.size()) == 2:
-            n1, d = x1.size()
-            n2, d = x2.size()
-        else:
-            b, n1, d = x1.size()
-            _, n2, _ = x2.size()
+        batch_shape = x1.shape[:-2]
+        n1, d = x1.shape[-2:]
+        n2 = x2.shape[-2]
 
-        K = torch.zeros(b, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)  # batch x n1(d+1) x n2(d+1)
+        K = torch.zeros(*batch_shape, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)
 
         if not diag:
             # Scale the inputs by the lengthscale (for stability)
@@ -66,7 +62,7 @@ class RBFKernelGrad(RBFKernel):
             x2_ = x2.div(self.lengthscale)
 
             # Form all possible rank-1 products for the gradient and Hessian blocks
-            outer = x1_.view([b, n1, 1, d]) - x2_.view([b, 1, n2, d])
+            outer = x1_.view(*batch_shape, n1, 1, d) - x2_.view(*batch_shape, 1, n2, d)
             outer = outer / self.lengthscale
             outer = torch.transpose(outer, -1, -2).contiguous()
 
@@ -76,19 +72,19 @@ class RBFKernelGrad(RBFKernel):
             K[..., :n1, :n2] = K_11
 
             # 2) First gradient block
-            outer1 = outer.view([b, n1, n2 * d])
+            outer1 = outer.view(*batch_shape, n1, n2 * d)
             K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])
 
             # 3) Second gradient block
-            outer2 = outer.transpose(-1, -3).contiguous().view([b, n2, n1 * d])
+            outer2 = outer.transpose(-1, -3).contiguous().view(*batch_shape, n2, n1 * d)
             outer2 = outer2.transpose(-1, -2)
             K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])
 
             # 4) Hessian block
             outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])
             kp = KroneckerProductLazyTensor(
-                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1]) / self.lengthscale.pow_(2),
-                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1])
+                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow_(2),
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1)
             )
             chain_rule = kp.evaluate() - outer3
             K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])
@@ -109,8 +105,8 @@ class RBFKernelGrad(RBFKernel):
                 raise RuntimeError("diag=True only works when x1 == x2")
 
             kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)
-            grad_diag = torch.ones(b, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow_(2)
-            grad_diag = grad_diag.transpose(-1, -2).contiguous().view([b, n2 * d])
+            grad_diag = torch.ones(*batch_shape, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow_(2)
+            grad_diag = grad_diag.transpose(-1, -2).contiguous().view(*batch_shape, n2 * d)
             k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)
             pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))
             return k_diag[..., pi]
diff --git a/gpytorch/kernels/scale_kernel.py b/gpytorch/kernels/scale_kernel.py
index ce004508..e235495b 100644
--- a/gpytorch/kernels/scale_kernel.py
+++ b/gpytorch/kernels/scale_kernel.py
@@ -31,7 +31,7 @@ class ScaleKernel(Kernel):
             The base kernel to be scaled.
         :attr:`batch_shape` (int, optional):
             Set this if you want a separate outputscale for each batch of input data. It should be `b`
-            if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`
+            if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([])`
         :attr:`outputscale_prior` (Prior, optional): Set this if you want to apply a prior to the outputscale
             parameter.  Default: `None`
         :attr:`outputscale_constraint` (Constraint, optional): Set this if you want to apply a constraint to the
@@ -56,7 +56,8 @@ class ScaleKernel(Kernel):
             outputscale_constraint = Positive()
 
         self.base_kernel = base_kernel
-        self.register_parameter(name="raw_outputscale", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape)))
+        outputscale = torch.zeros(*self.batch_shape) if len(self.batch_shape) else torch.tensor(0.)
+        self.register_parameter(name="raw_outputscale", parameter=torch.nn.Parameter(outputscale))
         if outputscale_prior is not None:
             self.register_prior(
                 "outputscale_prior", outputscale_prior, lambda: self.outputscale, lambda v: self._set_outputscale(v)
@@ -79,11 +80,11 @@ class ScaleKernel(Kernel):
 
     def forward(self, x1, x2, batch_dims=None, diag=False, **params):
         outputscales = self.outputscale
-        if batch_dims == (0, 2) and outputscales.numel() > 1:
-            outputscales = outputscales.unsqueeze(1).repeat(1, x1.size(-1)).view(-1)
+        if batch_dims == (0, 2):
+            outputscales = outputscales.unsqueeze(0).repeat(x1.size(-1), *([1] * outputscales.dim()))
 
         orig_output = self.base_kernel.forward(x1, x2, diag=diag, batch_dims=batch_dims, **params)
-        outputscales = outputscales.view(-1, *([1] * (orig_output.dim() - 1)))
+        outputscales = outputscales.view(*outputscales.shape, *([1] * (orig_output.dim() - outputscales.dim())))
 
         if diag:
             return delazify(orig_output) * outputscales
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
index 3649a5dd..ad63602f 100644
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -67,7 +67,7 @@ class SpectralMixtureKernel(Kernel):
         self,
         num_mixtures=None,
         ard_num_dims=1,
-        batch_shape=torch.Size([1]),
+        batch_shape=torch.Size([]),
         mixture_scales_prior=None,
         mixture_scales_constraint=None,
         mixture_means_prior=None,
@@ -81,9 +81,6 @@ class SpectralMixtureKernel(Kernel):
         if mixture_means_prior is not None or mixture_scales_prior is not None or mixture_weights_prior is not None:
             logger.warning("Priors not implemented for SpectralMixtureKernel")
 
-        if len(batch_shape) > 1:
-            raise RuntimeError("SpectralMixtureKernel does not yet support multiple batch dimensions.")
-
         # This kernel does not use the default lengthscale
         super(SpectralMixtureKernel, self).__init__(ard_num_dims=ard_num_dims, batch_shape=batch_shape, **kwargs)
         self.num_mixtures = num_mixtures
@@ -198,25 +195,18 @@ class SpectralMixtureKernel(Kernel):
         """
         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:])
+            x1_ = x1_.unsqueeze(0).transpose(0, -1)
             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:])
+                x2_ = x2_.unsqueeze(0).transpose(0, -1)
 
         if diag:
             return x1_, x2_
         else:
             return x1_.unsqueeze(-2), x2_.unsqueeze(-3)
 
-    def forward(self, x1, x2, **params):
-        if x1.dim() > 3 or x2.dim() > 3:
-            raise RuntimeError("SpectralMixtureKernel does not yet support multiple batch dimensions.")
-
+    def forward(self, x1, x2, batch_dims=None, **params):
         batch_shape = x1.shape[:-2]
         n, num_dims = x1.shape[-2:]
 
@@ -233,8 +223,8 @@ class SpectralMixtureKernel(Kernel):
 
         # Expand x1 and x2 to account for the number of mixtures
         # Should make x1/x2 (b x k x n x d) for k mixtures
-        x1_ = x1.unsqueeze(1)
-        x2_ = x2.unsqueeze(1)
+        x1_ = x1.unsqueeze(len(batch_shape))
+        x2_ = x2.unsqueeze(len(batch_shape))
 
         # Compute distances - scaled by appropriate parameters
         x1_exp = x1_ * self.mixture_scales
@@ -243,20 +233,31 @@ class SpectralMixtureKernel(Kernel):
         x2_cos = x2_ * self.mixture_means
 
         # Create grids
-        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)
+        x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, batch_dims=batch_dims, **params)
+        x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, batch_dims=batch_dims, **params)
 
         # Compute the exponential and cosine terms
         exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi ** 2)
         cos_term = (x1_cos_ - x2_cos_).mul_(2 * math.pi)
         res = exp_term.exp_() * cos_term.cos_()
 
-        # Product omer dimensions
-        res = res.prod(-1)
+        # Product over dimensions
+        if batch_dims == (0, 2):
+            res = res.squeeze(-1)
+        else:
+            res = res.prod(-1)
 
         # Sum over mixtures
         mixture_weights = self.mixture_weights
-        while mixture_weights.dim() < res.dim():
-            mixture_weights = mixture_weights.unsqueeze(-1)
-        res = (res * mixture_weights).sum(1)
+
+        if batch_dims == (0, 2):
+            while mixture_weights.dim() < res.dim() - 1:
+                mixture_weights = mixture_weights.unsqueeze(-1)
+            mixture_weights = mixture_weights.unsqueeze(0)
+            res = (res * mixture_weights).sum(len(batch_shape) + 1)
+        else:
+            while mixture_weights.dim() < res.dim():
+                mixture_weights = mixture_weights.unsqueeze(-1)
+
+            res = (res * mixture_weights).sum(len(batch_shape))
         return res
diff --git a/gpytorch/kernels/white_noise_kernel.py b/gpytorch/kernels/white_noise_kernel.py
index 66838242..f1644f8d 100644
--- a/gpytorch/kernels/white_noise_kernel.py
+++ b/gpytorch/kernels/white_noise_kernel.py
@@ -35,25 +35,12 @@ class WhiteNoiseKernel(Kernel):
         super(WhiteNoiseKernel, self).__init__()
         self.register_buffer("variances", variances)
 
-        # Enforce that variances are batch_size x data_size x task_size
-        # where batch_size and task_size are both 1 in the simplest case.
-        if variances.ndimension() == 3:
-            self.is_batch = True
-            self.variances = variances
-        elif variances.ndimension() == 2:
-            # By default, assume 2 dimensional data means data_size x task_size, and add a batch_size of 1
-            self.is_batch = False
-            self.variances = variances.unsqueeze(0)
-        else:
-            self.is_batch = False
-            self.variances = variances.unsqueeze(0).unsqueeze(-1)
-
     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)
-            return DiagLazyTensor(self.variances.view(self.variances.size(0), -1))
-        elif x1.size(-2) == x2.size(-2) and x1.size(-2) == self.variances.size(1) and torch.equal(x1, x2):
-            return DiagLazyTensor(self.variances.view(self.variances.size(0), -1))
+            return DiagLazyTensor(self.variances.view(*x1.shape[:-2], -1))
+        elif x1.size(-2) == x2.size(-2) and x1.size(-2) == self.variances.size(-1) and torch.equal(x1, x2):
+            return DiagLazyTensor(self.variances.view(*x1.shape[:-2], -1))
         else:
-            return ZeroLazyTensor(x1.size(-3), x1.size(-2), x2.size(-2), dtype=x1.dtype, device=x1.device)
+            return ZeroLazyTensor(*x1.shape[:-2], x1.size(-2), x2.size(-2), dtype=x1.dtype, device=x1.device)
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
index da62b024..152a63f0 100644
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -44,12 +44,44 @@ class LazyEvaluatedKernelTensor(LazyTensor):
         if self.batch_dims == (0, 2):
             x1 = x1.permute(0, 2, 1).contiguous()
             x1 = x1.view(-1, x1.size(-1), 1)
-        x1 = x1[(*batch_indices, row_index, _noop_index)]
+        try:
+            x1 = x1[(*batch_indices, row_index, _noop_index)]
+        # We're going to handle multi-batch indexing with a try-catch loop
+        # This way - in the default case, we can avoid doing expansions of x1 which can be timely
+        except IndexError:
+            if isinstance(batch_indices, slice):
+                x1 = x1.expand(1, *self.x1.shape[-2:])[(*batch_indices, row_index, _noop_index)]
+            elif isinstance(batch_indices, tuple):
+                if any([not isinstance(bi, slice) for bi in batch_indices]):
+                    raise RuntimeError(
+                        f"Attempting to tensor index a non-batch matrix's batch dimensions. "
+                        "Got batch index {batch_indices} but my shape was {self.shape}"
+                    )
+                x1 = x1.expand(
+                    *([1] * len(batch_indices)),
+                    *self.x1.shape[-2:]
+                )[(*batch_indices, row_index, _noop_index)]
         x2 = self.x2
         if self.batch_dims == (0, 2):
             x2 = x2.permute(0, 2, 1).contiguous()
             x2 = x2.view(-1, x2.size(-1), 1)
-        x2 = x2[(*batch_indices, col_index, _noop_index)]
+        try:
+            x2 = x2[(*batch_indices, col_index, _noop_index)]
+        # We're going to handle multi-batch indexing with a try-catch loop
+        # This way - in the default case, we can avoid doing expansions of x1 which can be timely
+        except IndexError:
+            if isinstance(batch_indices, slice):
+                x2 = x2.expand(1, *self.x2.shape[-2:])[(*batch_indices, row_index, _noop_index)]
+            elif isinstance(batch_indices, tuple):
+                if any([not isinstance(bi, slice) for bi in batch_indices]):
+                    raise RuntimeError(
+                        f"Attempting to tensor index a non-batch matrix's batch dimensions. "
+                        "Got batch index {batch_indices} but my shape was {self.shape}"
+                    )
+                x2 = x2.expand(
+                    *([1] * len(batch_indices)),
+                    *self.x2.shape[-2:]
+                )[(*batch_indices, row_index, _noop_index)]
 
         return self.__class__(
             x1, x2, kernel=self.kernel, **self.params
@@ -110,13 +142,11 @@ class LazyEvaluatedKernelTensor(LazyTensor):
 
         return x1.grad, x2.grad
 
+    @cached(name="size")
     def _size(self):
         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 torch.Size((self.x1.size(-1), ) + size)
         return size
 
     def _transpose_nonbatch(self):
@@ -143,37 +173,14 @@ class LazyEvaluatedKernelTensor(LazyTensor):
         x1 = self.x1
         x2 = self.x2
 
-        # TODO: until kernels support multi-batch mode, we have to ensure that the kernel has a batch dimension
-        is_batch = True
-        if x1.dim() == 2:
-            is_batch = False
-            x1 = x1.unsqueeze(0)
-            x2 = x2.unsqueeze(0)
-
         res = super(Kernel, self.kernel).__call__(x1, x2, diag=True, batch_dims=self.batch_dims, **self.params)
 
-        # TODO: until kernels support multi-batch mode, we have to remove any artificial batch dimension
-        # that we've added
-        if not is_batch:
-            x1 = x1.squeeze(0)
-            x2 = x2.squeeze(0)
-            res = res.squeeze(0)
-
-        # 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]))
+                expected_shape = torch.Size((x1.size(-1),) + shape[:-1])
                 if res.shape != expected_shape:
                     raise RuntimeError(
                         "The kernel {} is not equipped to handle batch_dims=(0, 2) "
diff --git a/gpytorch/utils/broadcasting.py b/gpytorch/utils/broadcasting.py
index c1826e79..5e11ae2a 100644
--- a/gpytorch/utils/broadcasting.py
+++ b/gpytorch/utils/broadcasting.py
@@ -3,7 +3,7 @@
 import torch
 
 
-def _mul_broadcast_shape(*shapes):
+def _mul_broadcast_shape(*shapes, error_msg=None):
     """Compute dimension suggested by multiple tensor indices (supports broadcasting)"""
 
     # Pad each shape so they have the same number of dimensions
@@ -16,7 +16,10 @@ def _mul_broadcast_shape(*shapes):
         non_singleton_sizes = tuple(size for size in size_by_dim if size != 1)
         if len(non_singleton_sizes):
             if any(size != non_singleton_sizes[0] for size in non_singleton_sizes):
-                raise RuntimeError("Shapes {} are not broadcastable for mul operation")
+                if error_msg is None:
+                    raise RuntimeError("Shapes are not broadcastable for mul operation")
+                else:
+                    raise RuntimeError(error_msg)
             final_size.append(non_singleton_sizes[0])
         # In this case - all dimensions are singleton sizes
         else:


Test patch:
diff --git a/test/examples/test_batch_gp_regression.py b/test/examples/test_batch_gp_regression.py
index 5639c709..b204a500 100644
--- a/test/examples/test_batch_gp_regression.py
+++ b/test/examples/test_batch_gp_regression.py
@@ -38,16 +38,26 @@ class ExactGPModel(gpytorch.models.ExactGP):
     def __init__(self, train_inputs, train_targets, likelihood, batch_size=1):
         super(ExactGPModel, self).__init__(train_inputs, train_targets, likelihood)
         self.mean_module = ConstantMean(batch_size=batch_size, prior=gpytorch.priors.SmoothedBoxPrior(-1, 1))
-        self.covar_module = ScaleKernel(
-            RBFKernel(
+        if batch_size > 1:
+            self.covar_module = ScaleKernel(
+                RBFKernel(
+                    batch_size=batch_size,
+                    lengthscale_prior=gpytorch.priors.NormalPrior(
+                        loc=torch.zeros(batch_size, 1, 1), scale=torch.ones(batch_size, 1, 1)
+                    ),
+                ),
                 batch_size=batch_size,
-                lengthscale_prior=gpytorch.priors.NormalPrior(
-                    loc=torch.zeros(batch_size, 1, 1), scale=torch.ones(batch_size, 1, 1)
+                outputscale_prior=gpytorch.priors.SmoothedBoxPrior(-2, 2),
+            )
+        else:
+            self.covar_module = ScaleKernel(
+                RBFKernel(
+                    lengthscale_prior=gpytorch.priors.NormalPrior(
+                        loc=torch.zeros(batch_size, 1, 1), scale=torch.ones(batch_size, 1, 1)
+                    ),
                 ),
-            ),
-            batch_size=batch_size,
-            outputscale_prior=gpytorch.priors.SmoothedBoxPrior(-2, 2),
-        )
+                outputscale_prior=gpytorch.priors.SmoothedBoxPrior(-2, 2),
+            )
 
     def forward(self, x):
         mean_x = self.mean_module(x)
@@ -184,9 +194,9 @@ class TestBatchGPRegression(unittest.TestCase):
     def test_train_on_batch_test_on_batch_shared_hypers_over_batch(self):
         # We're manually going to set the hyperparameters to something they shouldn't be
         likelihood = GaussianLikelihood(
-            noise_prior=gpytorch.priors.NormalPrior(loc=torch.zeros(2), scale=torch.ones(2)), batch_size=1
+            noise_prior=gpytorch.priors.NormalPrior(loc=torch.zeros(2), scale=torch.ones(2))
         )
-        gp_model = ExactGPModel(train_x12, train_y12, likelihood, batch_size=1)
+        gp_model = ExactGPModel(train_x12, train_y12, likelihood)
         mll = gpytorch.ExactMarginalLogLikelihood(likelihood, gp_model)
 
         # Find optimal model hyperparameters
diff --git a/test/examples/test_batch_multitask_gp_regression.py b/test/examples/test_batch_multitask_gp_regression.py
index 89223bcf..18e2ff86 100644
--- a/test/examples/test_batch_multitask_gp_regression.py
+++ b/test/examples/test_batch_multitask_gp_regression.py
@@ -37,16 +37,27 @@ class ExactGPModel(gpytorch.models.ExactGP):
         self.mean_module = MultitaskMean(
             ConstantMean(batch_size=batch_size, prior=gpytorch.priors.SmoothedBoxPrior(-1, 1)), num_tasks=2
         )
-        self.covar_module = MultitaskKernel(
-            RBFKernel(
-                batch_size=batch_size,
-                lengthscale_prior=gpytorch.priors.NormalPrior(
-                    loc=torch.zeros(batch_size, 1, 1), scale=torch.ones(batch_size, 1, 1)
+        if batch_size > 1:
+            self.covar_module = MultitaskKernel(
+                RBFKernel(
+                    batch_size=batch_size,
+                    lengthscale_prior=gpytorch.priors.NormalPrior(
+                        loc=torch.zeros(batch_size, 1, 1), scale=torch.ones(batch_size, 1, 1)
+                    ),
                 ),
-            ),
-            num_tasks=2,
-            rank=1,
-        )
+                num_tasks=2,
+                rank=1,
+            )
+        else:
+            self.covar_module = MultitaskKernel(
+                RBFKernel(
+                    lengthscale_prior=gpytorch.priors.NormalPrior(
+                        loc=torch.zeros(batch_size, 1, 1), scale=torch.ones(batch_size, 1, 1)
+                    ),
+                ),
+                num_tasks=2,
+                rank=1,
+            )
 
     def forward(self, x):
         mean_x = self.mean_module(x)
diff --git a/test/examples/test_grid_gp_regression.py b/test/examples/test_grid_gp_regression.py
index 609fe98f..3f107129 100644
--- a/test/examples/test_grid_gp_regression.py
+++ b/test/examples/test_grid_gp_regression.py
@@ -80,7 +80,7 @@ class TestGridGPRegression(unittest.TestCase):
 
         optimizer = optim.Adam(list(gp_model.parameters()) + list(likelihood.parameters()), lr=0.1)
         optimizer.n_iter = 0
-        with gpytorch.settings.debug(False):
+        with gpytorch.settings.debug(True):
             for _ in range(20):
                 optimizer.zero_grad()
                 output = gp_model(train_x)
@@ -89,7 +89,7 @@ class TestGridGPRegression(unittest.TestCase):
                 optimizer.n_iter += 1
                 optimizer.step()
 
-            for param in gp_model.parameters():
+            for name, param in gp_model.named_parameters():
                 self.assertTrue(param.grad is not None)
                 self.assertGreater(param.grad.norm().item(), 0)
             for param in likelihood.parameters():
diff --git a/test/kernels/test_cosine_kernel.py b/test/kernels/test_cosine_kernel.py
index e422ddab..ed79c914 100644
--- a/test/kernels/test_cosine_kernel.py
+++ b/test/kernels/test_cosine_kernel.py
@@ -79,18 +79,18 @@ class TestCosineKernel(unittest.TestCase):
         self.assertLess(torch.norm(res - actual), 1e-5)
 
         # batch_dims
-        actual = torch.zeros(4, 3, 3)
+        actual = torch.zeros(2, 2, 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]))
+                        actual[l, k, 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))])
+        actual = actual.diagonal(dim1=-2, dim2=-1)
         self.assertLess(torch.norm(res - actual), 1e-5)
 
 
diff --git a/test/kernels/test_linear_kernel.py b/test/kernels/test_linear_kernel.py
index 1932aec8..ffa66bce 100644
--- a/test/kernels/test_linear_kernel.py
+++ b/test/kernels/test_linear_kernel.py
@@ -67,14 +67,14 @@ class TestLinearKernel(unittest.TestCase, BaseKernelTestCase):
 
         # batch_dims
         dim_group_a = a
-        dim_group_a = dim_group_a.permute(0, 2, 1).contiguous().view(-1, 3)
-        actual = torch.mul(dim_group_a.unsqueeze(-1), dim_group_a.unsqueeze(-2))
+        dim_group_a = dim_group_a.unsqueeze(0).transpose(0, -1)
+        actual = dim_group_a.matmul(dim_group_a.transpose(-2, -1))
         res = kernel(a, a, batch_dims=(0, 2)).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-4)
 
         # 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))])
+        actual = actual.diagonal(dim1=-2, dim2=-1)
         self.assertLess(torch.norm(res - actual), 1e-4)
 
 
diff --git a/test/kernels/test_matern_kernel.py b/test/kernels/test_matern_kernel.py
index c82a7d17..aaedb108 100644
--- a/test/kernels/test_matern_kernel.py
+++ b/test/kernels/test_matern_kernel.py
@@ -136,23 +136,21 @@ class TestMaternKernel(unittest.TestCase, BaseKernelTestCase):
         # 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,
+                [[[0.0, 0.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]],
+                [[[1.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
+                [[[0.0, 0.0], [0.0, 0.0]], [[4.0, 4.0], [0.0, 0.0]]],
+            ]
         )
+
         dist.mul_(math.sqrt(5))
+        dist = dist.view(3, 2, 2, 2)
         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))])
+        actual = actual.diagonal(dim1=-2, dim2=-1)
         self.assertLess(torch.norm(res - actual), 1e-5)
 
 
diff --git a/test/kernels/test_rbf_kernel.py b/test/kernels/test_rbf_kernel.py
index 637468ae..48418db3 100644
--- a/test/kernels/test_rbf_kernel.py
+++ b/test/kernels/test_rbf_kernel.py
@@ -15,9 +15,9 @@ class TestRBFKernel(unittest.TestCase, BaseKernelTestCase):
         return RBFKernel(ard_num_dims=num_dims, **kwargs)
 
     def test_ard(self):
-        a = torch.tensor([[[1, 2], [2, 4]]], dtype=torch.float)
-        b = torch.tensor([[[1, 3], [0, 4]]], dtype=torch.float)
-        lengthscales = torch.tensor([1, 2], dtype=torch.float).view(1, 1, 2)
+        a = torch.tensor([[1, 2], [2, 4]], dtype=torch.float)
+        b = torch.tensor([[1, 3], [0, 4]], dtype=torch.float)
+        lengthscales = torch.tensor([1, 2], dtype=torch.float).view(1, 2)
 
         kernel = RBFKernel(ard_num_dims=2)
         kernel.initialize(lengthscale=lengthscales)
@@ -31,7 +31,7 @@ class TestRBFKernel(unittest.TestCase, BaseKernelTestCase):
 
         # Diag
         res = kernel(a, b).diag()
-        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        actual = actual.diag()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
         # batch_dims
@@ -66,14 +66,16 @@ class TestRBFKernel(unittest.TestCase, BaseKernelTestCase):
         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)
+        double_batch_a = scaled_a.unsqueeze(0).transpose(0, -1)
+        double_batch_b = scaled_b.unsqueeze(0).transpose(0, -1)
+        actual = double_batch_a.transpose(-1, -2).unsqueeze(-1) - double_batch_b.transpose(-1, -2).unsqueeze(-2)
+        actual = actual.pow(2).mul_(-0.5).exp().view(3, 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))])
+        actual = actual.diagonal(dim1=-2, dim2=-1)
         self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_ard_separate_batch(self):
diff --git a/test/kernels/test_scale_kernel.py b/test/kernels/test_scale_kernel.py
index adcc5caa..8608c12d 100644
--- a/test/kernels/test_scale_kernel.py
+++ b/test/kernels/test_scale_kernel.py
@@ -7,9 +7,9 @@ from gpytorch.kernels import RBFKernel, ScaleKernel
 
 class TestScaleKernel(unittest.TestCase):
     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)
+        a = torch.tensor([[1, 2], [2, 4]], dtype=torch.float)
+        b = torch.tensor([[1, 3], [0, 4]], dtype=torch.float)
+        lengthscales = torch.tensor([1, 2], dtype=torch.float).view(1, 2)
 
         base_kernel = RBFKernel(ard_num_dims=2)
         base_kernel.initialize(lengthscale=lengthscales)
@@ -26,12 +26,12 @@ class TestScaleKernel(unittest.TestCase):
 
         # Diag
         res = kernel(a, b).diag()
-        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
+        actual = actual.diag()
         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 = actual.pow(2).mul_(-0.5).exp().view(2, 2, 2)
         actual.mul_(3)
         res = kernel(a, b, batch_dims=(0, 2)).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
@@ -65,15 +65,17 @@ class TestScaleKernel(unittest.TestCase):
         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)
+        double_batch_a = scaled_a.unsqueeze(0).transpose(0, -1)
+        double_batch_b = scaled_b.unsqueeze(0).transpose(0, -1)
+        actual = double_batch_a.transpose(-1, -2).unsqueeze(-1) - double_batch_b.transpose(-1, -2).unsqueeze(-2)
+        actual = actual.pow(2).mul_(-0.5).exp().view(3, 2, 2, 2)
+        actual[:, 1, :, :].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))])
+        actual = actual.diagonal(dim1=-2, dim2=-1)
         self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_initialize_outputscale(self):
diff --git a/test/kernels/test_spectral_mixture_kernel.py b/test/kernels/test_spectral_mixture_kernel.py
index 4de8f555..07448e5b 100644
--- a/test/kernels/test_spectral_mixture_kernel.py
+++ b/test/kernels/test_spectral_mixture_kernel.py
@@ -11,12 +11,14 @@ class TestSpectralMixtureKernel(unittest.TestCase):
         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]
+        scales = scales.unsqueeze(1)
+        means = means.unsqueeze(1)
+        weights = torch.tensor([4, 2], dtype=torch.float)
         kernel = SpectralMixtureKernel(num_mixtures=2, ard_num_dims=2)
         kernel.initialize(
-            mixture_weights=torch.tensor([[4, 2]], dtype=torch.float),
-            mixture_means=torch.tensor([[[[1, 2]], [[2, 1]]]], dtype=torch.float),
-            mixture_scales=torch.tensor([[[[0.5, 0.25]], [[0.25, 0.5]]]], dtype=torch.float),
+            mixture_weights=weights,
+            mixture_means=means,
+            mixture_scales=scales,
         )
         kernel.eval()
 
diff --git a/test/kernels/test_white_noise_kernel.py b/test/kernels/test_white_noise_kernel.py
index acd73d83..057174d1 100644
--- a/test/kernels/test_white_noise_kernel.py
+++ b/test/kernels/test_white_noise_kernel.py
@@ -61,7 +61,7 @@ class TestWhiteNoiseKernel(unittest.TestCase):
 
     def test_computes_diag_eval_batch(self):
         a = torch.tensor([[4, 2, 8], [4, 2, 8]], dtype=torch.float).view(2, 3, 1)
-        variances = torch.randn(2, 3, 1)
+        variances = torch.randn(2, 3)
         kernel = WhiteNoiseKernel(variances=variances)
         kernel.eval()
         actual = torch.cat(