Repository: cornellius-gp/gpytorch
Pull request #589 (2019-03-20 23:29:14)

Support for multiple batch dimensions in *base* kernels
First of a few PRs to support multiple batch dimensions throughout GPyTorch. This specifically adds support to most of the base kernels.
- Simplifies kernel constructors to not repeat work done in the Kernel constructor. This includes collapsing unused args in to **kwargs. The argument for doing this is that the kernels were extremely inconsistent about which super args they explicitly took in the first place. If we want them expanded back out, there should at least be consistency.
- Nearly all base kernels + ScaleKernel now support multiple batch dimension inputs.
- `batch_size` kernel argument has been deprecated in favor of `batch_shape`, which takes a `torch.Size` and is used to make multibatch lengthscales etc.
- Added a battery of unit tests for batch and multibatch kernel calling to the `BaseKernelTestCase` class.

Todo:
- [x] Finish updating docstrings


Patch:
diff --git a/gpytorch/kernels/cosine_kernel.py b/gpytorch/kernels/cosine_kernel.py
index 5c756e3d..8c5771cb 100644
--- a/gpytorch/kernels/cosine_kernel.py
+++ b/gpytorch/kernels/cosine_kernel.py
@@ -3,7 +3,6 @@
 import math
 import torch
 from .kernel import Kernel
-from torch.nn.functional import softplus
 
 
 class CosineKernel(Kernel):
@@ -21,9 +20,9 @@ class CosineKernel(Kernel):
     where :math:`p` is the periord length parameter.
 
     Args:
-        :attr:`batch_size` (int, optional):
+        :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: `1`
+            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`
         :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`.
@@ -40,7 +39,7 @@ class CosineKernel(Kernel):
 
     Attributes:
         :attr:`period_length` (Tensor):
-            The period length parameter. Size = `batch_size x 1 x 1`.
+            The period length parameter. Size = `*batch_shape x 1 x 1`.
 
     Example:
         >>> x = torch.randn(10, 5)
@@ -51,25 +50,17 @@ class CosineKernel(Kernel):
         >>> # Batch: Simple option
         >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.CosineKernel())
         >>> # Batch: different lengthscale for each batch
-        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.CosineKernel(batch_size=2))
+        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.CosineKernel(batch_shape=torch.Size([2])))
         >>> covar = covar_module(x)  # Output: LazyVariable of size (2 x 10 x 10)
     """
 
-    def __init__(
-        self,
-        active_dims=None,
-        batch_size=1,
-        period_length_prior=None,
-        eps=1e-6,
-        param_transform=softplus,
-        inv_param_transform=None,
-        **kwargs
-    ):
-        super(CosineKernel, self).__init__(
-            active_dims=active_dims, param_transform=param_transform, inv_param_transform=inv_param_transform
+    def __init__(self, period_length_prior=None, **kwargs):
+        super(CosineKernel, self).__init__(**kwargs)
+
+        self.register_parameter(
+            name="raw_period_length", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, 1))
         )
-        self.eps = eps
-        self.register_parameter(name="raw_period_length", parameter=torch.nn.Parameter(torch.zeros(batch_size, 1, 1)))
+
         if period_length_prior is not None:
             self.register_prior(
                 "period_length_prior",
diff --git a/gpytorch/kernels/index_kernel.py b/gpytorch/kernels/index_kernel.py
index ad0f544e..64641ab9 100644
--- a/gpytorch/kernels/index_kernel.py
+++ b/gpytorch/kernels/index_kernel.py
@@ -3,7 +3,6 @@
 import torch
 from .kernel import Kernel
 from ..lazy import DiagLazyTensor, InterpolatedLazyTensor, PsdSumLazyTensor, RootLazyTensor
-from torch.nn.functional import softplus
 
 
 class IndexKernel(Kernel):
@@ -22,7 +21,7 @@ class IndexKernel(Kernel):
     Args:
         :attr:`num_tasks` (int):
             Total number of indices.
-        :attr:`batch_size` (int, optional):
+        :attr:`batch_shape` (torch.Size, optional):
             Set if the MultitaskKernel is operating on batches of data (and you want different
             parameters for each batch)
         :attr:`rank` (int):
@@ -42,14 +41,21 @@ class IndexKernel(Kernel):
             The element-wise log of the :math:`\mathbf v` vector.
     """
 
-    def __init__(self, num_tasks, rank=1, batch_size=1, prior=None, param_transform=softplus, inv_param_transform=None):
+    def __init__(
+        self,
+        num_tasks,
+        rank=1,
+        prior=None,
+        **kwargs
+    ):
         if rank > num_tasks:
             raise RuntimeError("Cannot create a task covariance matrix larger than the number of tasks")
-        super(IndexKernel, self).__init__(param_transform=param_transform, inv_param_transform=inv_param_transform)
+        super().__init__(**kwargs)
+
         self.register_parameter(
-            name="covar_factor", parameter=torch.nn.Parameter(torch.randn(batch_size, num_tasks, rank))
+            name="covar_factor", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks, rank))
         )
-        self.register_parameter(name="raw_var", parameter=torch.nn.Parameter(torch.randn(batch_size, num_tasks)))
+        self.register_parameter(name="raw_var", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks)))
         if prior is not None:
             self.register_prior("IndexKernelPrior", prior, self._eval_covar_matrix)
 
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 42bf54ee..346ea98a 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -7,7 +7,9 @@ from ..lazy import lazify, delazify, LazyEvaluatedKernelTensor, ZeroLazyTensor
 from ..module import Module
 from .. import settings
 from ..utils.transforms import _get_inv_param_transform
+from ..utils.deprecation import _deprecate_kwarg_with_transform
 from torch.nn.functional import softplus
+from ..utils import broadcasting
 
 
 @torch.jit.script
@@ -83,7 +85,7 @@ class Kernel(Module):
       This is controlled by the `ard_num_dims` keyword argument (as well has `has_lengthscale=True`).
 
     In batch-mode (i.e. when :math:`x_1` and :math:`x_2` are batches of input matrices), each
-    batch of data can have its own lengthscale parameter by setting the `batch_size`
+    batch of data can have its own lengthscale parameter by setting the `batch_shape`
     keyword argument to the appropriate number of batches.
 
     .. note::
@@ -97,9 +99,9 @@ class Kernel(Module):
         :attr:`ard_num_dims` (int, optional):
             Set this if you want a separate lengthscale for each input
             dimension. It should be `d` if :attr:`x1` is a `n x d` matrix.  Default: `None`
-        :attr:`batch_size` (int, optional):
+        :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: `1`
+            data. It should be `b1 x ... x bk` if :attr:`x1` is a `b1 x ... x bk x n x d` tensor.  Default: `1`
         :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`.
@@ -116,7 +118,7 @@ class Kernel(Module):
     Base Attributes:
         :attr:`lengthscale` (Tensor):
             The lengthscale parameter. Size/shape of parameter depends on the
-            :attr:`ard_num_dims` and :attr:`batch_size` arguments.
+            :attr:`ard_num_dims` and :attr:`batch_shape` arguments.
 
     Example:
         >>> covar_module = gpytorch.kernels.LinearKernel()
@@ -129,7 +131,7 @@ class Kernel(Module):
         self,
         has_lengthscale=False,
         ard_num_dims=None,
-        batch_size=1,
+        batch_shape=torch.Size([1]),
         active_dims=None,
         lengthscale_prior=None,
         param_transform=softplus,
@@ -142,15 +144,21 @@ class Kernel(Module):
             active_dims = torch.tensor(active_dims, dtype=torch.long)
         self.register_buffer("active_dims", active_dims)
         self.ard_num_dims = ard_num_dims
-        self.batch_size = batch_size
+
+        self.batch_shape = _deprecate_kwarg_with_transform(
+            kwargs, "batch_size", "batch_shape", batch_shape, lambda n: torch.Size([n])
+        )
+
         self.__has_lengthscale = has_lengthscale
         self._param_transform = param_transform
         self._inv_param_transform = _get_inv_param_transform(param_transform, inv_param_transform)
+        self.eps = eps
+
         if has_lengthscale:
-            self.eps = eps
             lengthscale_num_dims = 1 if ard_num_dims is None else ard_num_dims
             self.register_parameter(
-                name="raw_lengthscale", parameter=torch.nn.Parameter(torch.zeros(batch_size, 1, lengthscale_num_dims))
+                name="raw_lengthscale",
+                parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, lengthscale_num_dims))
             )
             if lengthscale_prior is not None:
                 self.register_prior(
@@ -194,11 +202,13 @@ class Kernel(Module):
         self.initialize(raw_lengthscale=self._inv_param_transform(value))
 
     def size(self, x1, x2):
-        non_batch_size = (x1.size(-2), x2.size(-2))
-        if x1.ndimension() == 3:
-            return torch.Size((x1.size(0),) + non_batch_size)
-        else:
-            return torch.Size(non_batch_size)
+        expected_size = broadcasting._matmul_broadcast_shape(
+            torch.Size([*x1.shape[:-2], x1.size(-2), x1.size(-1)]),
+            torch.Size([*x2.shape[:-2], x2.size(-1), x2.size(-2)]),
+            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))
+        )
+        return expected_size
 
     @abstractmethod
     def forward(self, x1, x2, diag=False, batch_dims=None, **params):
diff --git a/gpytorch/kernels/linear_kernel.py b/gpytorch/kernels/linear_kernel.py
index f4af0b51..f363c524 100644
--- a/gpytorch/kernels/linear_kernel.py
+++ b/gpytorch/kernels/linear_kernel.py
@@ -39,8 +39,8 @@ class LinearKernel(Kernel):
             `len(active_dims)` should equal `num_dimensions`.
     """
 
-    def __init__(self, num_dimensions=None, offset_prior=None, variance_prior=None, active_dims=None):
-        super(LinearKernel, self).__init__(active_dims=active_dims)
+    def __init__(self, num_dimensions=None, offset_prior=None, variance_prior=None, **kwargs):
+        super(LinearKernel, self).__init__(**kwargs)
         if num_dimensions is not None:
             warnings.warn(
                 "The `num_dimensions` argument is deprecated and no longer used.",
@@ -55,7 +55,9 @@ class LinearKernel(Kernel):
                 "The `offset_prior` argument is deprecated and no longer used.",
                 DeprecationWarning
             )
-        self.register_parameter(name="raw_variance", parameter=torch.nn.Parameter(torch.zeros(1)))
+        self.register_parameter(
+            name="raw_variance", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, 1))
+        )
         if variance_prior is not None:
             self.register_prior(
                 "variance_prior",
diff --git a/gpytorch/kernels/matern_kernel.py b/gpytorch/kernels/matern_kernel.py
index b838ea61..13180022 100644
--- a/gpytorch/kernels/matern_kernel.py
+++ b/gpytorch/kernels/matern_kernel.py
@@ -3,7 +3,6 @@
 import math
 import torch
 from .kernel import Kernel
-from torch.nn.functional import softplus
 from ..functions import MaternCovariance
 
 
@@ -41,9 +40,9 @@ class MaternKernel(Kernel):
         :attr:`ard_num_dims` (int, optional):
             Set this if you want a separate lengthscale for each
             input dimension. It should be `d` if :attr:`x1` is a `n x d` matrix. Default: `None`
-        :attr:`batch_size` (int, optional):
+        :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: `1`
+             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`
         :attr:`active_dims` (tuple of ints, optional):
             Set this if you want to
             compute the covariance of only a few input dimensions. The ints
@@ -62,7 +61,7 @@ class MaternKernel(Kernel):
     Attributes:
         :attr:`lengthscale` (Tensor):
             The lengthscale parameter. Size/shape of parameter depends on the
-            :attr:`ard_num_dims` and :attr:`batch_size` arguments.
+            :attr:`ard_num_dims` and :attr:`batch_shape` arguments.
 
     Example:
         >>> x = torch.randn(10, 5)
@@ -76,34 +75,14 @@ class MaternKernel(Kernel):
         >>> # Batch: Simple option
         >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.MaternKernel(nu=0.5))
         >>> # Batch: different lengthscale for each batch
-        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.MaternKernel(nu=0.5, batch_size=2))
+        >>> covar_module = gpytorch.kernels.MaternKernel(nu=0.5, batch_shape=torch.Size([2])
         >>> covar = covar_module(x)  # Output: LazyVariable of size (2 x 10 x 10)
     """
 
-    def __init__(
-        self,
-        nu=2.5,
-        ard_num_dims=None,
-        batch_size=1,
-        active_dims=None,
-        lengthscale_prior=None,
-        param_transform=softplus,
-        inv_param_transform=None,
-        eps=1e-6,
-        **kwargs
-    ):
+    def __init__(self, nu=2.5, **kwargs):
         if nu not in {0.5, 1.5, 2.5}:
             raise RuntimeError("nu expected to be 0.5, 1.5, or 2.5")
-        super(MaternKernel, self).__init__(
-            has_lengthscale=True,
-            ard_num_dims=ard_num_dims,
-            batch_size=batch_size,
-            active_dims=active_dims,
-            lengthscale_prior=lengthscale_prior,
-            param_transform=param_transform,
-            inv_param_transform=inv_param_transform,
-            eps=eps,
-        )
+        super(MaternKernel, self).__init__(has_lengthscale=True, **kwargs)
         self.nu = nu
 
     def forward(self, x1, x2, **params):
diff --git a/gpytorch/kernels/multitask_kernel.py b/gpytorch/kernels/multitask_kernel.py
index 394c5a98..1cabd2cd 100644
--- a/gpytorch/kernels/multitask_kernel.py
+++ b/gpytorch/kernels/multitask_kernel.py
@@ -29,16 +29,15 @@ class MultitaskKernel(Kernel):
             Prior to use for task kernel. See :class:`gpytorch.kernels.IndexKernel` for details.
     """
 
-    def __init__(self, data_covar_module, num_tasks, rank=1, batch_size=1, task_covar_prior=None):
+    def __init__(self, data_covar_module, num_tasks, rank=1, task_covar_prior=None, **kwargs):
         """
         """
-        super(MultitaskKernel, self).__init__()
+        super(MultitaskKernel, self).__init__(**kwargs)
         self.task_covar_module = IndexKernel(
-            num_tasks=num_tasks, batch_size=batch_size, rank=rank, prior=task_covar_prior
+            num_tasks=num_tasks, batch_shape=self.batch_shape, rank=rank, prior=task_covar_prior
         )
         self.data_covar_module = data_covar_module
         self.num_tasks = num_tasks
-        self.batch_size = 1
 
     def forward(self, x1, x2, diag=False, batch_dims=None, **params):
         if batch_dims == (0, 2):
diff --git a/gpytorch/kernels/periodic_kernel.py b/gpytorch/kernels/periodic_kernel.py
index 723d8ea9..0e00a8e0 100644
--- a/gpytorch/kernels/periodic_kernel.py
+++ b/gpytorch/kernels/periodic_kernel.py
@@ -3,7 +3,6 @@
 import math
 import torch
 from .kernel import Kernel
-from torch.nn.functional import softplus
 
 
 class PeriodicKernel(Kernel):
@@ -33,9 +32,9 @@ class PeriodicKernel(Kernel):
         This kernel does not have an ARD lengthscale option.
 
     Args:
-        :attr:`batch_size` (int, optional):
+        :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: `1`.
+             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.
         :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`.
@@ -54,9 +53,9 @@ class PeriodicKernel(Kernel):
 
     Attributes:
         :attr:`lengthscale` (Tensor):
-            The lengthscale parameter. Size = `batch_size x 1 x 1`.
+            The lengthscale parameter. Size = `*batch_shape x 1 x 1`.
         :attr:`period_length` (Tensor):
-            The period length parameter. Size = `batch_size x 1 x 1`.
+            The period length parameter. Size = `*batch_shape x 1 x 1`.
 
     Example:
         >>> x = torch.randn(10, 5)
@@ -71,27 +70,12 @@ class PeriodicKernel(Kernel):
         >>> covar = covar_module(x)  # Output: LazyVariable of size (2 x 10 x 10)
     """
 
-    def __init__(
-        self,
-        active_dims=None,
-        batch_size=1,
-        lengthscale_prior=None,
-        period_length_prior=None,
-        param_transform=softplus,
-        inv_param_transform=None,
-        eps=1e-6,
-        **kwargs
-    ):
-        super(PeriodicKernel, self).__init__(
-            has_lengthscale=True,
-            active_dims=active_dims,
-            batch_size=batch_size,
-            lengthscale_prior=lengthscale_prior,
-            param_transform=param_transform,
-            inv_param_transform=inv_param_transform,
-            eps=eps,
-        )
-        self.register_parameter(name="raw_period_length", parameter=torch.nn.Parameter(torch.zeros(batch_size, 1, 1)))
+    def __init__(self, period_length_prior=None, **kwargs):
+        super(PeriodicKernel, self).__init__(has_lengthscale=True, **kwargs)
+        self.register_parameter(
+            name="raw_period_length",
+            parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, 1)))
+
         if period_length_prior is not None:
             self.register_prior(
                 "period_length_prior",
diff --git a/gpytorch/kernels/rbf_kernel.py b/gpytorch/kernels/rbf_kernel.py
index ce5dc083..32d8388d 100644
--- a/gpytorch/kernels/rbf_kernel.py
+++ b/gpytorch/kernels/rbf_kernel.py
@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 from .kernel import Kernel
-from torch.nn.functional import softplus
 import torch
 from ..functions import RBFCovariance
 
@@ -35,9 +34,9 @@ class RBFKernel(Kernel):
         :attr:`ard_num_dims` (int, optional):
             Set this if you want a separate lengthscale for each
             input dimension. It should be `d` if :attr:`x1` is a `n x d` matrix. Default: `None`
-        :attr:`batch_size` (int, optional):
+        :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: `1`.
+            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.
         :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`.
@@ -54,7 +53,7 @@ class RBFKernel(Kernel):
     Attributes:
         :attr:`lengthscale` (Tensor):
             The lengthscale parameter. Size/shape of parameter depends on the
-            :attr:`ard_num_dims` and :attr:`batch_size` arguments.
+            :attr:`ard_num_dims` and :attr:`batch_shape` arguments.
 
     Example:
         >>> x = torch.randn(10, 5)
@@ -68,31 +67,12 @@ class RBFKernel(Kernel):
         >>> # Batch: Simple option
         >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
         >>> # Batch: different lengthscale for each batch
-        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(batch_size=2))
+        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(batch_shape=torch.Size([2])))
         >>> covar = covar_module(x)  # Output: LazyTensor of size (2 x 10 x 10)
     """
 
-    def __init__(
-        self,
-        ard_num_dims=None,
-        batch_size=1,
-        active_dims=None,
-        lengthscale_prior=None,
-        param_transform=softplus,
-        inv_param_transform=None,
-        eps=1e-6,
-        **kwargs
-    ):
-        super(RBFKernel, self).__init__(
-            has_lengthscale=True,
-            ard_num_dims=ard_num_dims,
-            batch_size=batch_size,
-            active_dims=active_dims,
-            lengthscale_prior=lengthscale_prior,
-            param_transform=param_transform,
-            inv_param_transform=inv_param_transform,
-            eps=eps,
-        )
+    def __init__(self, **kwargs):
+        super(RBFKernel, self).__init__(has_lengthscale=True, **kwargs)
 
     def forward(self, x1, x2, diag=False, **params):
         if (
diff --git a/gpytorch/kernels/rbf_kernel_grad.py b/gpytorch/kernels/rbf_kernel_grad.py
index d873980c..195faf45 100644
--- a/gpytorch/kernels/rbf_kernel_grad.py
+++ b/gpytorch/kernels/rbf_kernel_grad.py
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 from .rbf_kernel import RBFKernel
-from torch.nn.functional import softplus
 import torch
 from ..lazy.kronecker_product_lazy_tensor import KroneckerProductLazyTensor
 
@@ -19,9 +18,9 @@ class RBFKernelGrad(RBFKernel):
         decorate this kernel with a :class:`gpytorch.kernels.ScaleKernel`.
 
     Args:
-        :attr:`batch_size` (int, optional):
+        :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: `1`.
+             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.
         :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`.
@@ -38,7 +37,7 @@ class RBFKernelGrad(RBFKernel):
     Attributes:
         :attr:`lengthscale` (Tensor):
             The lengthscale parameter. Size/shape of parameter depends on the
-            :attr:`ard_num_dims` and :attr:`batch_size` arguments.
+            :attr:`ard_num_dims` and :attr:`batch_shape` arguments.
 
     Example:
         >>> x = torch.randn(10, 5)
@@ -50,31 +49,9 @@ class RBFKernelGrad(RBFKernel):
         >>> # Batch: Simple option
         >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad())
         >>> # Batch: different lengthscale for each batch
-        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad(batch_size=2))
+        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad(batch_shape=torch.Size([2])))
         >>> covar = covar_module(x)  # Output: LazyTensor of size (2 x 60 x 60)
     """
-
-    def __init__(
-        self,
-        batch_size=1,
-        active_dims=None,
-        lengthscale_prior=None,
-        param_transform=softplus,
-        inv_param_transform=None,
-        eps=1e-6,
-        **kwargs
-    ):
-        # TODO: Add support for ARD
-        super(RBFKernelGrad, self).__init__(
-            batch_size=batch_size,
-            active_dims=active_dims,
-            lengthscale_prior=lengthscale_prior,
-            param_transform=softplus,
-            inv_param_transform=inv_param_transform,
-            eps=eps,
-            **kwargs
-        )
-
     def forward(self, x1, x2, diag=False, **params):
         b = 1
         if len(x1.size()) == 2:
diff --git a/gpytorch/kernels/scale_kernel.py b/gpytorch/kernels/scale_kernel.py
index 84981a57..f13858b6 100644
--- a/gpytorch/kernels/scale_kernel.py
+++ b/gpytorch/kernels/scale_kernel.py
@@ -2,8 +2,6 @@
 
 import torch
 from .kernel import Kernel
-from ..utils.transforms import _get_inv_param_transform
-from torch.nn.functional import softplus
 from ..lazy import delazify
 
 
@@ -20,7 +18,7 @@ class ScaleKernel(Kernel):
     where :math:`\theta_\text{scale}` is the `outputscale` parameter.
 
     In batch-mode (i.e. when :math:`x_1` and :math:`x_2` are batches of input matrices), each
-    batch of data can have its own `outputscale` parameter by setting the `batch_size`
+    batch of data can have its own `outputscale` parameter by setting the `batch_shape`
     keyword argument to the appropriate number of batches.
 
     .. note::
@@ -30,9 +28,9 @@ class ScaleKernel(Kernel):
     Args:
         :attr:`base_kernel` (Kernel):
             The base kernel to be scaled.
-        :attr:`batch_size` (int, optional):
+        :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: `1`
+            if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`
         :attr:`outputscale_prior` (Prior, optional): Set this if you want to apply a prior to the outputscale
             parameter.  Default: `None`
         :attr:`param_transform` (function, optional):
@@ -45,7 +43,7 @@ class ScaleKernel(Kernel):
         :attr:`base_kernel` (Kernel):
             The kernel module to be scaled.
         :attr:`outputscale` (Tensor):
-            The outputscale parameter. Size/shape of parameter depends on the :attr:`batch_size` arguments.
+            The outputscale parameter. Size/shape of parameter depends on the :attr:`batch_shape` arguments.
 
     Example:
         >>> x = torch.randn(10, 5)
@@ -54,20 +52,10 @@ class ScaleKernel(Kernel):
         >>> covar = scaled_covar_module(x)  # Output: LazyTensor of size (10 x 10)
     """
 
-    def __init__(
-        self,
-        base_kernel,
-        batch_size=1,
-        outputscale_prior=None,
-        param_transform=softplus,
-        inv_param_transform=None,
-        **kwargs
-    ):
-        super(ScaleKernel, self).__init__(has_lengthscale=False, batch_size=batch_size)
+    def __init__(self, base_kernel, outputscale_prior=None, **kwargs):
+        super(ScaleKernel, self).__init__(has_lengthscale=False, **kwargs)
         self.base_kernel = base_kernel
-        self._param_transform = param_transform
-        self._inv_param_transform = _get_inv_param_transform(param_transform, inv_param_transform)
-        self.register_parameter(name="raw_outputscale", parameter=torch.nn.Parameter(torch.zeros(batch_size)))
+        self.register_parameter(name="raw_outputscale", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape)))
         if outputscale_prior is not None:
             self.register_prior(
                 "outputscale_prior", outputscale_prior, lambda: self.outputscale, lambda v: self._set_outputscale(v)
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
index 3d97cdf1..0e34432b 100644
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -4,7 +4,6 @@ import logging
 import math
 import torch
 from .kernel import Kernel
-from torch.nn.functional import softplus
 
 logger = logging.getLogger()
 
@@ -19,7 +18,8 @@ class SpectralMixtureKernel(Kernel):
 
         Unlike other kernels,
         * :attr:`ard_num_dums` **must equal** the number of dimensions of the data
-        * :attr:`batch_size` **must equal** the batch size of the data (1 if the data is not batched)
+        * :attr:`batch_shape` **must equal** the batch size of the data (torch.Size([1]) if the data is not batched)
+        * :attr:`batch_shape` **cannot** contain more than one batch dimension.
         * This kernel should not be combined with a :class:`gpytorch.kernels.ScaleKernel`.
 
     Args:
@@ -28,9 +28,9 @@ class SpectralMixtureKernel(Kernel):
         :attr:`ard_num_dims` (int, optional):
             Set this to match the dimensionality of the input.
             It should be `d` if :attr:`x1` is a `n x d` matrix. Default: `1`
-        :attr:`batch_size` (int, optional):
+        :attr:`batch_shape` (torch.Size, optional):
             Set this if the data is
-            batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `1`
+             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`
         :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`.
@@ -66,20 +66,14 @@ class SpectralMixtureKernel(Kernel):
     .. _Gaussian Process Kernels for Pattern Discovery and Extrapolation:
         https://arxiv.org/pdf/1302.4245.pdf
     """
-    # TODO: add equation to docs
-
     def __init__(
         self,
         num_mixtures=None,
         ard_num_dims=1,
-        batch_size=1,
-        active_dims=None,
-        eps=1e-6,
+        batch_shape=torch.Size([1]),
         mixture_scales_prior=None,
         mixture_means_prior=None,
         mixture_weights_prior=None,
-        param_transform=softplus,
-        inv_param_transform=None,
         **kwargs
     ):
         if num_mixtures is None:
@@ -87,19 +81,17 @@ 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__(
-            active_dims=active_dims, param_transform=param_transform, inv_param_transform=inv_param_transform
-        )
+        super(SpectralMixtureKernel, self).__init__(ard_num_dims=ard_num_dims, batch_shape=batch_shape, **kwargs)
         self.num_mixtures = num_mixtures
-        self.batch_size = batch_size
-        self.ard_num_dims = ard_num_dims
-        self.eps = eps
 
         self.register_parameter(
-            name="raw_mixture_weights", parameter=torch.nn.Parameter(torch.zeros(self.batch_size, self.num_mixtures))
+            name="raw_mixture_weights", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, self.num_mixtures))
         )
-        ms_shape = torch.Size([self.batch_size, self.num_mixtures, 1, self.ard_num_dims])
+        ms_shape = torch.Size([*self.batch_shape, self.num_mixtures, 1, self.ard_num_dims])
         self.register_parameter(name="raw_mixture_means", parameter=torch.nn.Parameter(torch.zeros(ms_shape)))
         self.register_parameter(name="raw_mixture_scales", parameter=torch.nn.Parameter(torch.zeros(ms_shape)))
 
@@ -205,18 +197,21 @@ class SpectralMixtureKernel(Kernel):
             return x1_.unsqueeze(-2), x2_.unsqueeze(-3)
 
     def forward(self, x1, x2, **params):
-        batch_size, n, num_dims = x1.size()
-        _, m, _ = x2.size()
+        if x1.dim() > 3 or x2.dim() > 3:
+            raise RuntimeError("SpectralMixtureKernel does not yet support multiple batch dimensions.")
+
+        batch_shape = x1.shape[:-2]
+        n, num_dims = x1.shape[-2:]
 
         if not num_dims == self.ard_num_dims:
             raise RuntimeError(
                 "The SpectralMixtureKernel expected the input to have {} dimensionality "
                 "(based on the ard_num_dims argument). Got {}.".format(self.ard_num_dims, num_dims)
             )
-        if not batch_size == self.batch_size:
+        if not batch_shape == self.batch_shape:
             raise RuntimeError(
                 "The SpectralMixtureKernel expected the input to have a batch_size of {} "
-                "(based on the batch_size argument). Got {}.".format(self.batch_size, batch_size)
+                "(based on the batch_size argument). Got {}.".format(self.batch_shape, batch_shape)
             )
 
         # Expand x1 and x2 to account for the number of mixtures
diff --git a/gpytorch/utils/broadcasting.py b/gpytorch/utils/broadcasting.py
index 911d8230..361e1507 100644
--- a/gpytorch/utils/broadcasting.py
+++ b/gpytorch/utils/broadcasting.py
@@ -25,17 +25,23 @@ def _mul_broadcast_shape(*shapes):
     return torch.Size(final_size)
 
 
-def _matmul_broadcast_shape(shape_a, shape_b):
+def _matmul_broadcast_shape(shape_a, shape_b, error_msg=None):
     """Compute dimension of matmul operation on shapes (supports broadcasting)"""
     m, n, p = shape_a[-2], shape_a[-1], shape_b[-1]
 
     if len(shape_b) == 1:
         if n != p:
-            raise RuntimeError("Incompatible dimensions for matmul")
+            if error_msg is None:
+                raise RuntimeError("Incompatible dimensions for matmul")
+            else:
+                raise RuntimeError(error_msg)
         return shape_a[:-1]
 
     if n != shape_b[-2]:
-        raise RuntimeError("Incompatible dimensions for matmul")
+        if error_msg is not None:
+            raise RuntimeError("Incompatible dimensions for matmul")
+        else:
+            raise RuntimeError(error_msg)
 
     tail_shape = torch.Size([m, p])
     bc_shape = _mul_broadcast_shape(shape_a[:-2], shape_b[:-2])
diff --git a/gpytorch/utils/deprecation.py b/gpytorch/utils/deprecation.py
index faf3c3dd..39e8e368 100644
--- a/gpytorch/utils/deprecation.py
+++ b/gpytorch/utils/deprecation.py
@@ -30,6 +30,14 @@ def _deprecate_kwarg(kwargs, old_kw, new_kw, new_kw_value):
     return new_kw_value
 
 
+def _deprecate_kwarg_with_transform(kwargs, old_kw, new_kw, new_kw_value, transform):
+    old_kwarg = kwargs.get(old_kw)
+    if old_kwarg is not None:
+        warnings.warn("The `{}` argument is deprecated. Use `{}` instead.".format(old_kw, new_kw), DeprecationWarning)
+        return transform(old_kwarg)
+    return new_kw_value
+
+
 def _deprecated_renamed_method(cls, old_method_name, new_method_name):
     def _deprecated_method(self, *args, **kwargs):
         warnings.warn(


Test patch:
diff --git a/test/kernels/_base_kernel_test_case.py b/test/kernels/_base_kernel_test_case.py
index f307628a..a7fbe9d4 100644
--- a/test/kernels/_base_kernel_test_case.py
+++ b/test/kernels/_base_kernel_test_case.py
@@ -9,6 +9,10 @@ class BaseKernelTestCase(object):
     def create_kernel_no_ard(self, **kwargs):
         raise NotImplementedError()
 
+    @abstractmethod
+    def create_kernel_ard(self, num_dims, **kwargs):
+        raise NotImplementedError()
+
     def test_active_dims_list(self):
         kernel = self.create_kernel_no_ard(active_dims=[0, 2, 4, 6])
         x = torch.randn(50, 10)
@@ -16,7 +20,7 @@ class BaseKernelTestCase(object):
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().evaluate()
 
-        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-5)
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-4)
 
     def test_active_dims_range(self):
         active_dims = list(range(3, 9))
@@ -26,4 +30,78 @@ class BaseKernelTestCase(object):
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().evaluate()
 
-        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-5)
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-4)
+
+    def test_no_batch_kernel_single_batch_x_no_ard(self):
+        kernel = self.create_kernel_no_ard()
+        x = torch.randn(2, 50, 2)
+        batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+
+        actual_mat_1 = kernel(x[0]).evaluate_kernel().evaluate()
+        actual_mat_2 = kernel(x[1]).evaluate_kernel().evaluate()
+        actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
+
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+
+    def test_single_batch_kernel_single_batch_x_no_ard(self):
+        kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
+        x = torch.randn(2, 50, 2)
+        batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+
+        actual_mat_1 = kernel(x[0]).evaluate_kernel().evaluate()
+        actual_mat_2 = kernel(x[1]).evaluate_kernel().evaluate()
+        actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
+
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+
+    def test_no_batch_kernel_double_batch_x_no_ard(self):
+        kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
+        x = torch.randn(3, 2, 50, 2)
+        batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+
+        ij_actual_covars = []
+        for i in range(x.size(0)):
+            i_actual_covars = []
+            for j in range(x.size(1)):
+                i_actual_covars.append(kernel(x[i, j]).evaluate_kernel().evaluate())
+            ij_actual_covars.append(torch.cat([ac.unsqueeze(0) for ac in i_actual_covars]))
+
+        actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
+
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+
+    def test_no_batch_kernel_double_batch_x_ard(self):
+        try:
+            kernel = self.create_kernel_ard(num_dims=2, batch_shape=torch.Size([]))
+        except NotImplementedError:
+            return
+
+        x = torch.randn(3, 2, 50, 2)
+        batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+
+        ij_actual_covars = []
+        for i in range(x.size(0)):
+            i_actual_covars = []
+            for j in range(x.size(1)):
+                i_actual_covars.append(kernel(x[i, j]).evaluate_kernel().evaluate())
+            ij_actual_covars.append(torch.cat([ac.unsqueeze(0) for ac in i_actual_covars]))
+
+        actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
+
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+
+    def test_smoke_double_batch_kernel_double_batch_x_no_ard(self):
+        kernel = self.create_kernel_no_ard(batch_shape=torch.Size([3, 2]))
+        x = torch.randn(3, 2, 50, 2)
+        batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+        return batch_covar_mat
+
+    def test_smoke_double_batch_kernel_double_batch_x_ard(self):
+        try:
+            kernel = self.create_kernel_ard(num_dims=2, batch_shape=torch.Size([3, 2]))
+        except NotImplementedError:
+            return
+
+        x = torch.randn(3, 2, 50, 2)
+        batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+        return batch_covar_mat
diff --git a/test/kernels/test_linear_kernel.py b/test/kernels/test_linear_kernel.py
index c7effe18..1932aec8 100644
--- a/test/kernels/test_linear_kernel.py
+++ b/test/kernels/test_linear_kernel.py
@@ -8,7 +8,7 @@ from test.kernels._base_kernel_test_case import BaseKernelTestCase
 
 class TestLinearKernel(unittest.TestCase, BaseKernelTestCase):
     def create_kernel_no_ard(self, **kwargs):
-        return LinearKernel(num_dimensions=10, **kwargs)
+        return LinearKernel(**kwargs)
 
     def test_computes_linear_function_rectangular(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
diff --git a/test/kernels/test_matern_kernel.py b/test/kernels/test_matern_kernel.py
index 93724a99..c82a7d17 100644
--- a/test/kernels/test_matern_kernel.py
+++ b/test/kernels/test_matern_kernel.py
@@ -7,10 +7,29 @@ from gpytorch.kernels import MaternKernel
 from test.kernels._base_kernel_test_case import BaseKernelTestCase
 
 
+class TestMatern25BaseKernel(unittest.TestCase, BaseKernelTestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        return MaternKernel(nu=2.5, **kwargs)
+
+    def create_kernel_ard(self, num_dims, **kwargs):
+        return MaternKernel(nu=2.5, ard_num_dims=num_dims, **kwargs)
+
+
+class TestMatern05BaseKernel(unittest.TestCase, BaseKernelTestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        return MaternKernel(nu=0.5, **kwargs)
+
+    def create_kernel_ard(self, num_dims, **kwargs):
+        return MaternKernel(nu=0.5, ard_num_dims=num_dims, **kwargs)
+
+
 class TestMaternKernel(unittest.TestCase, BaseKernelTestCase):
     def create_kernel_no_ard(self, **kwargs):
         return MaternKernel(nu=1.5, **kwargs)
 
+    def create_kernel_ard(self, num_dims, **kwargs):
+        return MaternKernel(nu=1.5, ard_num_dims=num_dims, **kwargs)
+
     def test_forward_nu_1_over_2(self):
         a = torch.tensor([4, 2, 8], dtype=torch.float).view(3, 1)
         b = torch.tensor([0, 2], dtype=torch.float).view(2, 1)
diff --git a/test/kernels/test_rbf_kernel.py b/test/kernels/test_rbf_kernel.py
index 4c386a9b..637468ae 100644
--- a/test/kernels/test_rbf_kernel.py
+++ b/test/kernels/test_rbf_kernel.py
@@ -11,6 +11,9 @@ class TestRBFKernel(unittest.TestCase, BaseKernelTestCase):
     def create_kernel_no_ard(self, **kwargs):
         return RBFKernel(**kwargs)
 
+    def create_kernel_ard(self, num_dims, **kwargs):
+        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)