Repository: cornellius-gp/gpytorch

Issue:
[Docs] How to use SpectralMixtureKernel for BatchIndependent MultiOutput GP
# 📚 Documentation/Examples

** Is there documentation missing? **
This is a question about a specific use case.

I'd like to use the SpectralMixtureKernel to map a common set of input data to multiple outputs, without needing to determine correlation between the different outputs. It seems this is what the `BatchIndependent` (https://docs.gpytorch.ai/en/v1.2.1/examples/03_Multitask_Exact_GPs/Batch_Independent_Multioutput_GP.html) framework is for. However, I am having a hard time getting it to output the correct dimensions for the covariance matrices i.e. the expected size of the covariance matrices obtained from passing the input data through the SpectralMixtureKernel should be `[3, 100, 100]` for the scenario described below.

As an example, I have input data of size `[100, 15]` (i.e. 100 samples and 15 dimensions) and input targets of size `[100, 3]`. I'd like effectively create 3 independent GPs (one for each target dimension) that use the same input data. From the examples in the documentation (https://docs.gpytorch.ai/en/v1.2.1/kernels.html#spectralmixturekernel), it seems like batched and non-batched data give the same output covariance of size `[100,100]` which is effectively just representing a single GP. 

** Is documentation wrong? **
No

** Is there a feature that needs some example code? **
<!-- Let us know if we're missing any example notebooks -->

** Think you know how to fix the docs? ** (If so, we'd love a pull request from you!)

- Link to [GPyTorch documentation](https://gpytorch.readthedocs.io)
- Link to [GPyTorch examples](https://github.com/cornellius-gp/gpytorch/tree/master/examples)



Discussion:
Here is a code snippet of what I tried

```python
import torch
import gpytorch as gp


class SpectralMixtureGP(gp.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(SpectralMixtureGP, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gp.means.ConstantMean(batch_shape=torch.Size([train_y.size()[-1]]))
        self.covar_module = gp.kernels.SpectralMixtureKernel(num_mixtures=4, ard_num_dims=train_x.size()[-1],
                                                             batch_shape=torch.Size([train_y.size()[-1]])).cuda()
        self.covar_module.initialize_from_data(train_x, train_y)

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        print("Mean_x size: {}, Covar_x size: {}".format(mean_x.size(), covar_x.size()))
        return gp.distributions.MultitaskMultivariateNormal.from_batch_mvn(
                gp.distributions.MultivariateNormal(mean_x, covar_x)
            )


# Define data
train_x = torch.randn(100, 15).cuda()
train_y = torch.randn(100, 3).cuda()

# Define likelihood
likelihood = gp.likelihoods.MultitaskGaussianLikelihood(num_tasks=train_y.size()[-1]).cuda()

# Initialize model
model = SpectralMixtureGP(train_x, train_y, likelihood)

# Train the model
model.train()
likelihood.train()

mll = gp.mlls.ExactMarginalLogLikelihood(likelihood, model)

# Create optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

for _ in range(100):
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    loss.backward()
    optimizer.step()
```
The size of the mean and covariance matrices seem to be what I need, however, I get the following error. I understand that the `batch_shape` argument for the `SpectralMixtureKernel` is not following the documentation of the API where my `train_x` is of size `[100,15]` instead of `[3,100,15]`. The behavior of the the `batch_shape` argument seems to differ from that of the other standard kernels. I'm not sure how to get the intended result.

```
Traceback (most recent call last):
  File "/home/somesh/imitation/src/imitation/sandbox/test.py", line 46, in <module>
    loss = -mll(output, train_y)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/module.py", line 28, in __call__
    outputs = self.forward(*inputs, **kwargs)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/mlls/exact_marginal_log_likelihood.py", line 62, in forward
    res = output.log_prob(target)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/distributions/multitask_multivariate_normal.py", line 202, in log_prob
    return super().log_prob(value.view(*value.shape[:-2], -1))
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/distributions/multivariate_normal.py", line 140, in log_prob
    inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 1022, in inv_quad_logdet
    cholesky = CholLazyTensor(TriangularLazyTensor(self.cholesky()))
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 750, in cholesky
    chol = self._cholesky(upper=False)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/utils/memoize.py", line 59, in g
    return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 407, in _cholesky
    evaluated_kern_mat = self.evaluate_kernel()
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 914, in evaluate_kernel
    return self.representation_tree()(*self.representation())
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 1296, in representation_tree
    return LazyTensorRepresentationTree(self)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor_representation_tree.py", line 13, in __init__
    representation_size = len(arg.representation())
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py", line 1284, in representation
    representation += list(arg.representation())
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 307, in representation
    return self.evaluate_kernel().representation()
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/utils/memoize.py", line 59, in g
    return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 274, in evaluate_kernel
    res = self.kernel(x1, x2, diag=False, last_dim_is_batch=self.last_dim_is_batch, **self.params)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/kernels/kernel.py", line 396, in __call__
    res = lazify(super(Kernel, self).__call__(x1_, x2_, last_dim_is_batch=last_dim_is_batch, **params))
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/module.py", line 28, in __call__
    outputs = self.forward(*inputs, **kwargs)
  File "/home/somesh/imitation/lib/python3.7/site-packages/gpytorch/kernels/spectral_mixture_kernel.py", line 272, in forward
    "(based on the batch_size argument). Got {}.".format(self.batch_shape, batch_shape)
RuntimeError: The SpectralMixtureKernel expected the input to have a batch_size of torch.Size([3]) (based on the batch_size argument). Got torch.Size([]).
Mean_x size: torch.Size([3, 100]), Covar_x size: torch.Size([3, 100, 100])

Process finished with exit code 1
```

Resolving patch:
diff --git a/gpytorch/kernels/distributional_input_kernel.py b/gpytorch/kernels/distributional_input_kernel.py
--- a/gpytorch/kernels/distributional_input_kernel.py
+++ b/gpytorch/kernels/distributional_input_kernel.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+import torch
+
 from .kernel import Kernel
 
 
@@ -34,4 +36,7 @@ def forward(self, x1, x2, diag=False, *args, **kwargs):
         if not diag:
             return res
         else:
-            return res.diag()
+            if torch.is_tensor(res):
+                return res.diagonal(dim1=-1, dim2=-2)
+            else:
+                return res.diag()  # For LazyTensor
diff --git a/gpytorch/kernels/polynomial_kernel_grad.py b/gpytorch/kernels/polynomial_kernel_grad.py
--- a/gpytorch/kernels/polynomial_kernel_grad.py
+++ b/gpytorch/kernels/polynomial_kernel_grad.py
@@ -30,9 +30,17 @@ def forward(
             K22_base_diag = self.power * (self.power - 1) * base_diag.pow(self.power - 2)
             K12_base_diag = self.power * base_diag.pow(self.power - 1)
 
-            K22_diag = all_outers_diag * K22_base_diag.repeat(d) + K12_base_diag.repeat(d)
+            K22_diag = torch.add(
+                all_outers_diag * K22_base_diag.repeat(*([1] * (K22_base_diag.dim() - 1)), d),
+                K12_base_diag.repeat(*([1] * (K12_base_diag.dim() - 1)), d),
+            )
+
+            K_diag = torch.cat([K11_diag, K22_diag], dim=-1)
+            # Apply perfect shuffle
+            pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().reshape((n1 * (d + 1)))
+            K_diag = K_diag[..., pi1]
+            return K_diag
 
-            return torch.cat([K11_diag, K22_diag], dim=-1)
         else:
             base_inner_prod = torch.matmul(x1, x2.transpose(-2, -1)) + offset
             K11 = base_inner_prod.pow(self.power)
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -2,10 +2,12 @@
 
 import logging
 import math
+from typing import Tuple, Union
 
 import torch
 
-from ..constraints import Positive
+from ..constraints import Interval, Positive
+from ..priors import Prior
 from .kernel import Kernel
 
 logger = logging.getLogger()
@@ -14,42 +16,48 @@
 class SpectralMixtureKernel(Kernel):
     r"""
     Computes a covariance matrix based on the Spectral Mixture Kernel
-    between inputs :math:`\mathbf{x_1}` and :math:`\mathbf{x_2}`:
+    between inputs :math:`\mathbf{x_1}` and :math:`\mathbf{x_2}`.
+
     It was proposed in `Gaussian Process Kernels for Pattern Discovery and Extrapolation`_.
 
     .. note::
-
         Unlike other kernels,
-        * :attr:`ard_num_dims` **must equal** the number of dimensions of the data
-        * :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:
-        :attr:`num_mixtures` (int, optional):
-            The number of components in the mixture.
-        :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_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: `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`.
-        :attr:`eps` (float):
-            The minimum value that the lengthscale can take (prevents divide by zero errors). Default: `1e-6`.
-
-    Attributes:
-        :attr:`mixture_lengthscale` (Tensor):
-            The lengthscale parameter. Given `k` mixture components, and `b x n x d` data, this will be of
-            size `b x k x 1 x d`.
-        :attr:`mixture_means` (Tensor):
-            The mixture mean parameters (`b x k x 1 x d`).
-        :attr:`mixture_weights` (Tensor):
-            The mixture weight parameters (`b x k`).
+
+            * :attr:`ard_num_dims` **must equal** the number of dimensions of the data.
+            * This kernel should not be combined with a :class:`gpytorch.kernels.ScaleKernel`.
+
+    :param int num_mixtures: The number of components in the mixture.
+    :param int ard_num_dims: Set this to match the dimensionality of the input.
+        It should be `d` if :attr:`x1` is a `... x n x d` matrix. (Default: `1`.)
+    :param batch_shape: Set this if the data is batch of input data. It should
+        be `b_1 x ... x b_j` if :attr:`x1` is a `b_1 x ... x b_j x n x d` tensor. (Default: `torch.Size([])`.)
+    :type batch_shape: torch.Size, optional
+    :param active_dims: 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`.)
+    :type active_dims: float, optional
+    :param eps: The minimum value that the lengthscale can take (prevents divide by zero errors). (Default: `1e-6`.)
+    :type eps: float, optional
+
+    :param mixture_scales_prior: A prior to set on the :attr:`mixture_scales` parameter
+    :type mixture_scales_prior: ~gpytorch.priors.Prior, optional
+    :param mixture_scales_constraint: A constraint to set on the :attr:`mixture_scales` parameter
+    :type mixture_scales_constraint: ~gpytorch.constraints.Interval, optional
+    :param mixture_means_prior: A prior to set on the :attr:`mixture_means` parameter
+    :type mixture_means_prior: ~gpytorch.priors.Prior, optional
+    :param mixture_means_constraint: A constraint to set on the :attr:`mixture_means` parameter
+    :type mixture_means_constraint: ~gpytorch.constraints.Interval, optional
+    :param mixture_weights_prior: A prior to set on the :attr:`mixture_weights` parameter
+    :type mixture_weights_prior: ~gpytorch.priors.Prior, optional
+    :param mixture_weights_constraint: A constraint to set on the :attr:`mixture_weights` parameter
+    :type mixture_weights_constraint: ~gpytorch.constraints.Interval, optional
+
+    :ivar torch.Tensor mixture_scales: The lengthscale parameter. Given
+        `k` mixture components, and `... x n x d` data, this will be of size `... x k x 1 x d`.
+    :ivar torch.Tensor mixture_means: The mixture mean parameters (`... x k x 1 x d`).
+    :ivar torch.Tensor mixture_weights: The mixture weight parameters (`... x k`).
 
     Example:
+
         >>> # Non-batch
         >>> x = torch.randn(10, 5)
         >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, ard_num_dims=5)
@@ -60,7 +68,6 @@ class SpectralMixtureKernel(Kernel):
         >>> covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4, batch_size=2, ard_num_dims=5)
         >>> covar = covar_module(x)  # Output: LazyVariable of size (10 x 10)
 
-
     .. _Gaussian Process Kernels for Pattern Discovery and Extrapolation:
         https://arxiv.org/pdf/1302.4245.pdf
     """
@@ -69,15 +76,15 @@ class SpectralMixtureKernel(Kernel):
 
     def __init__(
         self,
-        num_mixtures=None,
-        ard_num_dims=1,
-        batch_shape=torch.Size([]),
-        mixture_scales_prior=None,
-        mixture_scales_constraint=None,
-        mixture_means_prior=None,
-        mixture_means_constraint=None,
-        mixture_weights_prior=None,
-        mixture_weights_constraint=None,
+        num_mixtures: int = None,
+        ard_num_dims: int = 1,
+        batch_shape: torch.Size = torch.Size([]),
+        mixture_scales_prior: Prior = None,
+        mixture_scales_constraint: Interval = None,
+        mixture_means_prior: Prior = None,
+        mixture_means_constraint: Interval = None,
+        mixture_weights_prior: Prior = None,
+        mixture_weights_constraint: Interval = None,
         **kwargs,
     ):
         if num_mixtures is None:
@@ -114,10 +121,10 @@ def mixture_scales(self):
         return self.raw_mixture_scales_constraint.transform(self.raw_mixture_scales)
 
     @mixture_scales.setter
-    def mixture_scales(self, value):
+    def mixture_scales(self, value: Union[torch.Tensor, float]):
         self._set_mixture_scales(value)
 
-    def _set_mixture_scales(self, value):
+    def _set_mixture_scales(self, value: Union[torch.Tensor, float]):
         if not torch.is_tensor(value):
             value = torch.as_tensor(value).to(self.raw_mixture_scales)
         self.initialize(raw_mixture_scales=self.raw_mixture_scales_constraint.inverse_transform(value))
@@ -127,10 +134,10 @@ def mixture_means(self):
         return self.raw_mixture_means_constraint.transform(self.raw_mixture_means)
 
     @mixture_means.setter
-    def mixture_means(self, value):
+    def mixture_means(self, value: Union[torch.Tensor, float]):
         self._set_mixture_means(value)
 
-    def _set_mixture_means(self, value):
+    def _set_mixture_means(self, value: Union[torch.Tensor, float]):
         if not torch.is_tensor(value):
             value = torch.as_tensor(value).to(self.raw_mixture_means)
         self.initialize(raw_mixture_means=self.raw_mixture_means_constraint.inverse_transform(value))
@@ -140,109 +147,148 @@ def mixture_weights(self):
         return self.raw_mixture_weights_constraint.transform(self.raw_mixture_weights)
 
     @mixture_weights.setter
-    def mixture_weights(self, value):
+    def mixture_weights(self, value: Union[torch.Tensor, float]):
         self._set_mixture_weights(value)
 
-    def _set_mixture_weights(self, value):
+    def _set_mixture_weights(self, value: Union[torch.Tensor, float]):
         if not torch.is_tensor(value):
             value = torch.as_tensor(value).to(self.raw_mixture_weights)
         self.initialize(raw_mixture_weights=self.raw_mixture_weights_constraint.inverse_transform(value))
 
-    def initialize_from_data_empspect(self, train_x, train_y):
+    def initialize_from_data_empspect(self, train_x: torch.Tensor, train_y: torch.Tensor):
         """
         Initialize mixture components based on the empirical spectrum of the data.
-
         This will often be better than the standard initialize_from_data method.
+
+        :param torch.Tensor train_x: Training inputs
+        :param torch.Tensor train_y: Training outputs
         """
+
         import numpy as np
         from scipy.fftpack import fft
         from scipy.integrate import cumtrapz
 
-        if not torch.is_tensor(train_x) or not torch.is_tensor(train_y):
-            raise RuntimeError("train_x and train_y should be tensors")
-        if train_x.ndimension() == 1:
-            train_x = train_x.unsqueeze(-1)
-
-        N = train_x.size(-2)
-        emp_spect = np.abs(fft(train_y.cpu().detach().numpy())) ** 2 / N
-        M = math.floor(N / 2)
-
-        freq1 = np.arange(M + 1)
-        freq2 = np.arange(-M + 1, 0)
-        freq = np.hstack((freq1, freq2)) / N
-        freq = freq[: M + 1]
-        emp_spect = emp_spect[: M + 1]
-
-        total_area = np.trapz(emp_spect, freq)
-        spec_cdf = np.hstack((np.zeros(1), cumtrapz(emp_spect, freq)))
-        spec_cdf = spec_cdf / total_area
-
-        a = np.random.rand(1000, self.ard_num_dims)
-        p, q = np.histogram(a, spec_cdf)
-        bins = np.digitize(a, q)
-        slopes = (spec_cdf[bins] - spec_cdf[bins - 1]) / (freq[bins] - freq[bins - 1])
-        intercepts = spec_cdf[bins - 1] - slopes * freq[bins - 1]
-        inv_spec = (a - intercepts) / slopes
-
-        from sklearn.mixture import GaussianMixture
-
-        GMM = GaussianMixture(n_components=self.num_mixtures, covariance_type="diag").fit(inv_spec)
-        means = GMM.means_
-        varz = GMM.covariances_
-        weights = GMM.weights_
-
-        self.mixture_means = means
-        self.mixture_scales = varz
-        self.mixture_weights = weights
-
-    def initialize_from_data(self, train_x, train_y, **kwargs):
-        if not torch.is_tensor(train_x) or not torch.is_tensor(train_y):
-            raise RuntimeError("train_x and train_y should be tensors")
-        if train_x.ndimension() == 1:
-            train_x = train_x.unsqueeze(-1)
-        if train_x.ndimension() == 2:
-            train_x = train_x.unsqueeze(0)
-
-        train_x_sort = train_x.sort(1)[0]
-        max_dist = train_x_sort[:, -1, :] - train_x_sort[:, 0, :]
-        min_dist_sort = (train_x_sort[:, 1:, :] - train_x_sort[:, :-1, :]).squeeze(0)
-        min_dist = torch.zeros(1, self.ard_num_dims, dtype=train_x.dtype, device=train_x.device)
-        for ind in range(self.ard_num_dims):
-            min_dist[:, ind] = min_dist_sort[((min_dist_sort[:, ind]).nonzero(as_tuple=False))[0], ind]
-
-        # Inverse of lengthscales should be drawn from truncated Gaussian | N(0, max_dist^2) |
-        self.raw_mixture_scales.data.normal_().mul_(max_dist).abs_().pow_(-1)
-        self.raw_mixture_scales.data = self.raw_mixture_scales_constraint.inverse_transform(
-            self.raw_mixture_scales.data
-        )
-        # Draw means from Unif(0, 0.5 / minimum distance between two points)
-        self.raw_mixture_means.data.uniform_().mul_(0.5).div_(min_dist)
-        self.raw_mixture_means.data = self.raw_mixture_means_constraint.inverse_transform(self.raw_mixture_means.data)
-        # Mixture weights should be roughly the stdv of the y values divided by the number of mixtures
-        self.raw_mixture_weights.data.fill_(train_y.std() / self.num_mixtures)
-        self.raw_mixture_weights.data = self.raw_mixture_weights_constraint.inverse_transform(
-            self.raw_mixture_weights.data
-        )
+        with torch.no_grad():
+            if not torch.is_tensor(train_x) or not torch.is_tensor(train_y):
+                raise RuntimeError("train_x and train_y should be tensors")
+            if train_x.ndimension() == 1:
+                train_x = train_x.unsqueeze(-1)
+            if self.active_dims is not None:
+                train_x = train_x[..., self.active_dims]
+
+            # Flatten batch dimensions
+            train_x = train_x.view(-1, train_x.size(-1))
+            train_y = train_y.view(-1)
+
+            N = train_x.size(-2)
+            emp_spect = np.abs(fft(train_y.cpu().detach().numpy())) ** 2 / N
+            M = math.floor(N / 2)
+
+            freq1 = np.arange(M + 1)
+            freq2 = np.arange(-M + 1, 0)
+            freq = np.hstack((freq1, freq2)) / N
+            freq = freq[: M + 1]
+            emp_spect = emp_spect[: M + 1]
+
+            total_area = np.trapz(emp_spect, freq)
+            spec_cdf = np.hstack((np.zeros(1), cumtrapz(emp_spect, freq)))
+            spec_cdf = spec_cdf / total_area
+
+            a = np.random.rand(1000, self.ard_num_dims)
+            p, q = np.histogram(a, spec_cdf)
+            bins = np.digitize(a, q)
+            slopes = (spec_cdf[bins] - spec_cdf[bins - 1]) / (freq[bins] - freq[bins - 1])
+            intercepts = spec_cdf[bins - 1] - slopes * freq[bins - 1]
+            inv_spec = (a - intercepts) / slopes
+
+            from sklearn.mixture import GaussianMixture
+
+            GMM = GaussianMixture(n_components=self.num_mixtures, covariance_type="diag").fit(inv_spec)
+            means = GMM.means_
+            varz = GMM.covariances_
+            weights = GMM.weights_
+
+            dtype = self.raw_mixture_means.dtype
+            device = self.raw_mixture_means.device
+            self.mixture_means = torch.tensor(means, dtype=dtype, device=device).unsqueeze(-2)
+            self.mixture_scales = torch.tensor(varz, dtype=dtype, device=device).unsqueeze(-2)
+            self.mixture_weights = torch.tensor(weights, dtype=dtype, device=device)
+
+    def initialize_from_data(self, train_x: torch.Tensor, train_y: torch.Tensor, **kwargs):
+        """
+        Initialize mixture components based on batch statistics of the data.
 
-    def _create_input_grid(self, x1, x2, diag=False, last_dim_is_batch=False, **params):
+        :param torch.Tensor train_x: Training inputs
+        :param torch.Tensor train_y: Training outputs
+        """
+
+        with torch.no_grad():
+            if not torch.is_tensor(train_x) or not torch.is_tensor(train_y):
+                raise RuntimeError("train_x and train_y should be tensors")
+            if train_x.ndimension() == 1:
+                train_x = train_x.unsqueeze(-1)
+            if self.active_dims is not None:
+                train_x = train_x[..., self.active_dims]
+
+            # Compute maximum distance between points in each dimension
+            train_x_sort = train_x.sort(dim=-2)[0]
+            max_dist = train_x_sort[..., -1, :] - train_x_sort[..., 0, :]
+
+            # Compute the minimum distance between points in each dimension
+            dists = train_x_sort[..., 1:, :] - train_x_sort[..., :-1, :]
+            # We don't want the minimum distance to be zero, so fill zero values with some large number
+            dists = torch.where(dists.eq(0.0), torch.tensor(1.0e10, dtype=train_x.dtype, device=train_x.device), dists)
+            sorted_dists = dists.sort(dim=-2)[0]
+            min_dist = sorted_dists[..., 0, :]
+
+            # Reshape min_dist and max_dist to match the shape of parameters
+            # First add a singleton data dimension (-2) and a dimension for the mixture components (-3)
+            min_dist = min_dist.unsqueeze_(-2).unsqueeze_(-3)
+            max_dist = max_dist.unsqueeze_(-2).unsqueeze_(-3)
+            # Compress any dimensions in min_dist/max_dist that correspond to singletons in the SM parameters
+            dim = -3
+            while -dim <= min_dist.dim():
+                if -dim > self.raw_mixture_scales.dim():
+                    min_dist = min_dist.min(dim=dim)[0]
+                    max_dist = max_dist.max(dim=dim)[0]
+                elif self.raw_mixture_scales.size(dim) == 1:
+                    min_dist = min_dist.min(dim=dim, keepdim=True)[0]
+                    max_dist = max_dist.max(dim=dim, keepdim=True)[0]
+                    dim -= 1
+                else:
+                    dim -= 1
+
+            # Inverse of lengthscales should be drawn from truncated Gaussian | N(0, max_dist^2) |
+            self.mixture_scales = torch.randn_like(self.raw_mixture_scales).mul_(max_dist).abs_().reciprocal_()
+            # Draw means from Unif(0, 0.5 / minimum distance between two points)
+            self.mixture_means = torch.rand_like(self.raw_mixture_means).mul_(0.5).div(min_dist)
+            # Mixture weights should be roughly the stdv of the y values divided by the number of mixtures
+            self.mixture_weights = train_y.std().div(self.num_mixtures)
+
+    def _create_input_grid(
+        self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, last_dim_is_batch: bool = False, **params
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         This is a helper method for creating a grid of the kernel's inputs.
         Use this helper rather than maually creating a meshgrid.
 
         The grid dimensions depend on the kernel's evaluation mode.
 
-        Args:
-            :attr:`x1` (Tensor `n x d` or `b x n x d`)
-            :attr:`x2` (Tensor `m x d` or `b x m x d`) - for diag mode, these must be the same inputs
-
-        Returns:
-            (:class:`Tensor`, :class:`Tensor) corresponding to the gridded `x1` and `x2`.
-            The shape depends on the kernel's mode
-
-            * `full_covar`: (`b x n x 1 x d` and `b x 1 x m x d`)
-            * `full_covar` with `last_dim_is_batch=True`: (`b x k x n x 1 x 1` and `b x k x 1 x m x 1`)
-            * `diag`: (`b x n x d` and `b x n x d`)
-            * `diag` with `last_dim_is_batch=True`: (`b x k x n x 1` and `b x k x n x 1`)
+        :param torch.Tensor x1: ... x n x d
+        :param torch.Tensor x2: ... x m x d (for diag mode, these must be the same inputs)
+        :param diag: Should the Kernel compute the whole kernel, or just the diag? (Default: True.)
+        :type diag: bool, optional
+        :param last_dim_is_batch: If this is true, it treats the last dimension
+            of the data as another batch dimension.  (Useful for additive
+            structure over the dimensions). (Default: False.)
+        :type last_dim_is_batch: bool, optional
+
+        :rtype: torch.Tensor, torch.Tensor
+        :return: Grid corresponding to x1 and x2. The shape depends on the kernel's mode:
+            * `full_covar`: (`... x n x 1 x d` and `... x 1 x m x d`)
+            * `full_covar` with `last_dim_is_batch=True`: (`... x k x n x 1 x 1` and `... x k x 1 x m x 1`)
+            * `diag`: (`... x n x d` and `... x n x d`)
+            * `diag` with `last_dim_is_batch=True`: (`... x k x n x 1` and `... x k x n x 1`)
         """
         x1_, x2_ = x1, x2
         if last_dim_is_batch:
@@ -257,8 +303,9 @@ def _create_input_grid(self, x1, x2, diag=False, last_dim_is_batch=False, **para
         else:
             return x1_.unsqueeze(-2), x2_.unsqueeze(-3)
 
-    def forward(self, x1, x2, last_dim_is_batch=False, **params):
-        batch_shape = x1.shape[:-2]
+    def forward(
+        self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, last_dim_is_batch: bool = False, **params
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         n, num_dims = x1.shape[-2:]
 
         if not num_dims == self.ard_num_dims:
@@ -266,16 +313,11 @@ def forward(self, x1, x2, last_dim_is_batch=False, **params):
                 "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_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_shape, batch_shape)
-            )
 
         # 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(len(batch_shape))
-        x2_ = x2.unsqueeze(len(batch_shape))
+        # Should make x1/x2 (... x k x n x d) for k mixtures
+        x1_ = x1.unsqueeze(-3)
+        x2_ = x2.unsqueeze(-3)
 
         # Compute distances - scaled by appropriate parameters
         x1_exp = x1_ * self.mixture_scales
@@ -284,8 +326,12 @@ def forward(self, x1, x2, last_dim_is_batch=False, **params):
         x2_cos = x2_ * self.mixture_means
 
         # Create grids
-        x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, last_dim_is_batch=last_dim_is_batch, **params)
-        x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, last_dim_is_batch=last_dim_is_batch, **params)
+        x1_exp_, x2_exp_ = self._create_input_grid(
+            x1_exp, x2_exp, diag=diag, last_dim_is_batch=last_dim_is_batch, **params
+        )
+        x1_cos_, x2_cos_ = self._create_input_grid(
+            x1_cos, x2_cos, diag=diag, last_dim_is_batch=last_dim_is_batch, **params
+        )
 
         # Compute the exponential and cosine terms
         exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi ** 2)
@@ -299,12 +345,11 @@ def forward(self, x1, x2, last_dim_is_batch=False, **params):
             res = res.prod(-1)
 
         # Sum over mixtures
-        mixture_weights = self.mixture_weights
-
-        if last_dim_is_batch:
+        mixture_weights = self.mixture_weights.unsqueeze(-1)
+        if not diag:
             mixture_weights = mixture_weights.unsqueeze(-1)
-        while mixture_weights.dim() < res.dim():
+        if last_dim_is_batch:
             mixture_weights = mixture_weights.unsqueeze(-1)
 
-        res = (res * mixture_weights).sum(len(batch_shape))
+        res = (res * mixture_weights).sum(-2 if diag else -3)
         return res


Test patch:
diff --git a/gpytorch/test/base_kernel_test_case.py b/gpytorch/test/base_kernel_test_case.py
--- a/gpytorch/test/base_kernel_test_case.py
+++ b/gpytorch/test/base_kernel_test_case.py
@@ -5,13 +5,14 @@
 
 import torch
 
+from .base_test_case import BaseTestCase
 
-class BaseKernelTestCase(object):
+
+class BaseKernelTestCase(BaseTestCase):
     @abstractmethod
     def create_kernel_no_ard(self, **kwargs):
         raise NotImplementedError()
 
-    @abstractmethod
     def create_kernel_ard(self, num_dims, **kwargs):
         raise NotImplementedError()
 
@@ -19,7 +20,7 @@ def create_data_no_batch(self):
         return torch.randn(50, 10)
 
     def create_data_single_batch(self):
-        return torch.randn(2, 50, 2)
+        return torch.randn(2, 3, 2)
 
     def create_data_double_batch(self):
         return torch.randn(3, 2, 50, 2)
@@ -31,7 +32,7 @@ def test_active_dims_list(self):
         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) / covar_mat_actual.norm(), 1e-4)
+        self.assertAllClose(covar_mat, covar_mat_actual, rtol=1e-3, atol=1e-5)
 
     def test_active_dims_range(self):
         active_dims = list(range(3, 9))
@@ -41,7 +42,7 @@ def test_active_dims_range(self):
         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) / covar_mat_actual.norm(), 1e-4)
+        self.assertAllClose(covar_mat, covar_mat_actual, rtol=1e-3, atol=1e-5)
 
     def test_no_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard()
@@ -52,7 +53,12 @@ def test_no_batch_kernel_single_batch_x_no_ard(self):
         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) / actual_covar_mat.norm(), 1e-4)
+        self.assertAllClose(batch_covar_mat, actual_covar_mat, rtol=1e-3, atol=1e-5)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
+        self.assertAllClose(kernel_diag, actual_diag, rtol=1e-3, atol=1e-5)
 
     def test_single_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
@@ -63,7 +69,12 @@ def test_single_batch_kernel_single_batch_x_no_ard(self):
         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) / actual_covar_mat.norm(), 1e-4)
+        self.assertAllClose(batch_covar_mat, actual_covar_mat, rtol=1e-3, atol=1e-5)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
+        self.assertAllClose(kernel_diag, actual_diag, rtol=1e-3, atol=1e-5)
 
     def test_no_batch_kernel_double_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
@@ -79,7 +90,12 @@ def test_no_batch_kernel_double_batch_x_no_ard(self):
 
         actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
+        self.assertAllClose(batch_covar_mat, actual_covar_mat, rtol=1e-3, atol=1e-5)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
+        self.assertAllClose(kernel_diag, actual_diag, rtol=1e-3, atol=1e-5)
 
     def test_no_batch_kernel_double_batch_x_ard(self):
         try:
@@ -99,12 +115,18 @@ def test_no_batch_kernel_double_batch_x_ard(self):
 
         actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 2e-4)
+        self.assertAllClose(batch_covar_mat, actual_covar_mat, rtol=1e-3, atol=1e-5)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
+        self.assertAllClose(kernel_diag, actual_diag, rtol=1e-3, atol=1e-5)
 
     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 = self.create_data_double_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+        kernel(x, diag=True)
         return batch_covar_mat
 
     def test_smoke_double_batch_kernel_double_batch_x_ard(self):
@@ -115,6 +137,7 @@ def test_smoke_double_batch_kernel_double_batch_x_ard(self):
 
         x = self.create_data_double_batch()
         batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+        kernel(x, diag=True)
         return batch_covar_mat
 
     def test_kernel_getitem_single_batch(self):
@@ -126,7 +149,7 @@ def test_kernel_getitem_single_batch(self):
         new_kernel = kernel[0]
         res2 = new_kernel(x[0]).evaluate()  # Should also be result of first kernel on first batch of data.
 
-        self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)
+        self.assertAllClose(res1, res2, rtol=1e-3, atol=1e-5)
 
     def test_kernel_getitem_double_batch(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([3, 2]))
@@ -137,7 +160,7 @@ def test_kernel_getitem_double_batch(self):
         new_kernel = kernel[0, 1]
         res2 = new_kernel(x[0, 1]).evaluate()  # Should also be result of first kernel on first batch of data.
 
-        self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)
+        self.assertAllClose(res1, res2, rtol=1e-3, atol=1e-5)
 
     def test_kernel_pickle_unpickle(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
diff --git a/test/examples/test_spectral_mixture_gp_regression.py b/test/examples/test_spectral_mixture_gp_regression.py
--- a/test/examples/test_spectral_mixture_gp_regression.py
+++ b/test/examples/test_spectral_mixture_gp_regression.py
@@ -72,7 +72,7 @@ def test_spectral_mixture_gp_mean_abs_error(self, empspect=False):
         # Optimize the model
         gp_model.train()
         likelihood.train()
-        optimizer = optim.SGD(list(gp_model.parameters()), lr=0.1)
+        optimizer = optim.Adam(list(gp_model.parameters()), lr=0.01)
         optimizer.n_iter = 0
 
         if not empspect:
diff --git a/test/kernels/test_matern_kernel.py b/test/kernels/test_matern_kernel.py
--- a/test/kernels/test_matern_kernel.py
+++ b/test/kernels/test_matern_kernel.py
@@ -19,10 +19,14 @@ def create_kernel_ard(self, num_dims, **kwargs):
 
 class TestMatern05BaseKernel(unittest.TestCase, BaseKernelTestCase):
     def create_kernel_no_ard(self, **kwargs):
-        return MaternKernel(nu=0.5, **kwargs)
+        kernel = MaternKernel(nu=0.5, **kwargs)
+        kernel.initialize(lengthscale=5.0)
+        return kernel
 
     def create_kernel_ard(self, num_dims, **kwargs):
-        return MaternKernel(nu=0.5, ard_num_dims=num_dims, **kwargs)
+        kernel = MaternKernel(nu=0.5, ard_num_dims=num_dims, **kwargs)
+        kernel.initialize(lengthscale=5.0)
+        return kernel
 
 
 class TestMaternKernel(unittest.TestCase, BaseKernelTestCase):
diff --git a/test/kernels/test_spectral_mixture_kernel.py b/test/kernels/test_spectral_mixture_kernel.py
--- a/test/kernels/test_spectral_mixture_kernel.py
+++ b/test/kernels/test_spectral_mixture_kernel.py
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 
-import math
 import unittest
 
 import torch
@@ -9,84 +8,131 @@
 
 
 class TestSpectralMixtureKernel(unittest.TestCase):
-    def test_standard(self):
-        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)
-        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=weights, mixture_means=means, mixture_scales=scales)
-        kernel.eval()
-
-        actual = torch.zeros(2, 3, 3, 2)
-        for i in range(3):
-            for j in range(3):
-                for k in range(2):
-                    new_term = torch.cos(2 * math.pi * (a[i] - a[j]) * means[k])
-                    new_term *= torch.exp(-2 * (math.pi * (a[i] - a[j])) ** 2 * scales[k] ** 2)
-                    actual[k, i, j] = new_term
-        actual = actual.prod(-1)
-        actual[0].mul_(weights[0])
-        actual[1].mul_(weights[1])
-        actual = actual.sum(0)
-
-        res = kernel(a, a).evaluate()
-        self.assertLess(torch.norm(res - actual), 1e-5)
-
-        # diag
-        res = kernel(a, a).diag()
-        actual = actual.diag()
-        self.assertLess(torch.norm(res - actual), 1e-5)
-
-        # batch_dims
-        actual = torch.zeros(2, 3, 3, 2)
-        for i in range(3):
-            for j in range(3):
-                for k in range(2):
-                    new_term = torch.cos(2 * math.pi * (a[i] - a[j]) * means[k])
-                    new_term *= torch.exp(-2 * (math.pi * (a[i] - a[j])) ** 2 * scales[k] ** 2)
-                    actual[k, i, j] = new_term
-        actual[0].mul_(weights[0])
-        actual[1].mul_(weights[1])
-        actual = actual.sum(0)
-        actual = actual.permute(2, 0, 1)
-        res = kernel(a, a, last_dim_is_batch=True).evaluate()
-        self.assertLess(torch.norm(res - actual), 1e-5)
-
-        # batch_dims + diag
-        res = kernel(a, a, last_dim_is_batch=True).diag()
-        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
-        self.assertLess(torch.norm(res - actual), 1e-5)
-
-    def test_batch_separate(self):
-        a = torch.tensor([[4, 2, 8], [1, 2, 3]], dtype=torch.float).view(2, 3, 1)
-        b = torch.tensor([[0, 2, 1], [-1, 2, 0]], dtype=torch.float).view(2, 3, 1)
-        means = torch.tensor([[1, 2], [2, 3]], dtype=torch.float).view(2, 2, 1, 1)
-        scales = torch.tensor([[0.5, 0.25], [0.25, 1]], dtype=torch.float).view(2, 2, 1, 1)
-        weights = torch.tensor([[4, 2], [1, 2]], dtype=torch.float).view(2, 2)
-        kernel = SpectralMixtureKernel(batch_shape=torch.Size([2]), num_mixtures=2)
-        kernel.initialize(mixture_weights=weights, mixture_means=means, mixture_scales=scales)
-        kernel.eval()
-
-        actual = torch.zeros(2, 3, 3)
-        for l in range(2):
-            for k in range(2):
-                for i in range(3):
-                    for j in range(3):
-                        new_term = torch.cos(2 * math.pi * (a[l, i] - b[l, j]) * means[l, k])
-                        new_term *= torch.exp(-2 * (math.pi * (a[l, i] - b[l, j])) ** 2 * scales[l, k] ** 2)
-                        new_term *= weights[l, k]
-                        actual[l, i, j] += new_term.item()
-
-        res = kernel(a, b).evaluate()
-        self.assertLess(torch.norm(res - actual), 1e-5)
-
-        # diag
-        res = kernel(a, b).diag()
-        actual = torch.cat([actual[i].diag().unsqueeze(0) for i in range(actual.size(0))])
-        self.assertLess(torch.norm(res - actual), 1e-5)
+    def create_kernel(self, num_dims, **kwargs):
+        return SpectralMixtureKernel(num_mixtures=5, ard_num_dims=num_dims, **kwargs)
+
+    def create_data_no_batch(self):
+        return torch.randn(50, 10)
+
+    def create_data_single_batch(self):
+        return torch.randn(2, 50, 2)
+
+    def create_data_double_batch(self):
+        return torch.randn(3, 2, 50, 2)
+
+    def test_active_dims_list(self):
+        x = self.create_data_no_batch()
+        kernel = self.create_kernel(num_dims=4, active_dims=[0, 2, 4, 6])
+        y = torch.randn_like(x[..., 0])
+        kernel.initialize_from_data(x, y)
+        kernel.initialize_from_data_empspect(x, y)
+
+        covar_mat = kernel(x).evaluate_kernel().evaluate()
+        kernel_basic = self.create_kernel(num_dims=4)
+        kernel_basic.raw_mixture_weights.data = kernel.raw_mixture_weights
+        kernel_basic.raw_mixture_means.data = kernel.raw_mixture_means
+        kernel_basic.raw_mixture_scales.data = kernel.raw_mixture_scales
+        covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().evaluate()
+
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
+
+    def test_active_dims_range(self):
+        active_dims = list(range(3, 9))
+        x = self.create_data_no_batch()
+        kernel = self.create_kernel(num_dims=6, active_dims=active_dims)
+        y = torch.randn_like(x[..., 0])
+        kernel.initialize_from_data(x, y)
+        kernel.initialize_from_data_empspect(x, y)
+
+        covar_mat = kernel(x).evaluate_kernel().evaluate()
+        kernel_basic = self.create_kernel(num_dims=6)
+        kernel_basic.raw_mixture_weights.data = kernel.raw_mixture_weights
+        kernel_basic.raw_mixture_means.data = kernel.raw_mixture_means
+        kernel_basic.raw_mixture_scales.data = kernel.raw_mixture_scales
+        covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().evaluate()
+
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
+
+    def test_no_batch_kernel_single_batch_x(self):
+        x = self.create_data_single_batch()
+        kernel = self.create_kernel(num_dims=x.size(-1))
+        y = torch.randn_like(x[..., 0])
+        kernel.initialize_from_data(x, y)
+        kernel.initialize_from_data_empspect(x, y)
+        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) / actual_covar_mat.norm(), 1e-4)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
+        self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
+
+    def test_single_batch_kernel_single_batch_x(self):
+        x = self.create_data_single_batch()
+        kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([]))
+        y = torch.randn_like(x[..., 0])
+        kernel.initialize_from_data(x, y)
+        kernel.initialize_from_data_empspect(x, y)
+        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) / actual_covar_mat.norm(), 1e-4)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
+        self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
+
+    def test_smoke_double_batch_kernel_double_batch_x(self):
+        x = self.create_data_double_batch()
+        kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([3, 2]))
+        y = torch.randn_like(x[..., 0])
+        kernel.initialize_from_data(x, y)
+        kernel.initialize_from_data_empspect(x, y)
+
+        batch_covar_mat = kernel(x).evaluate_kernel().evaluate()
+        kernel_diag = kernel(x, diag=True)
+        return batch_covar_mat, kernel_diag
+
+    def test_kernel_getitem_single_batch(self):
+        x = self.create_data_single_batch()
+        kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([2]))
+
+        res1 = kernel(x).evaluate()[0]  # Result of first kernel on first batch of data
+
+        new_kernel = kernel[0]
+        res2 = new_kernel(x[0]).evaluate()  # Should also be result of first kernel on first batch of data.
+
+        self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = res1.diagonal(dim1=-1, dim2=-2)
+        self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
+
+    def test_kernel_getitem_double_batch(self):
+        x = self.create_data_double_batch()
+        kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([3, 2]))
+
+        res1 = kernel(x).evaluate()[0, 1]  # Result of first kernel on first batch of data
+
+        new_kernel = kernel[0, 1]
+        res2 = new_kernel(x[0, 1]).evaluate()  # Should also be result of first kernel on first batch of data.
+
+        self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)
+
+        # Test diagonal
+        kernel_diag = kernel(x, diag=True)
+        actual_diag = res1.diagonal(dim1=-1, dim2=-2)
+        self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
 
 
 if __name__ == "__main__":