[Bug] gpytorch PeriodicKernel, diag() method yields to RuntimeError message

[LeoL135]

🐛 Bug

Hello,
I have a problem with the gpytorch Periodickernel.
When triying to call the diag() method, it yiels to a runtime error.
This leads to the problem that no confidence region or predictive variance can be computed later for predictions.
That happens both for few and large data.
I am not sure if this is a bug or am I doing something wrong here?

To reproduce

** Code snippet to reproduce **
covar_module = gpytorch.kernels.PeriodicKernel()
x1 = torch.randn(50)
lazy_covar_matrix = covar_module(x1) # Returns a RootLazyTensor
lazy_covar_matrix.diag()

** Stack trace/error message **

RuntimeError Traceback (most recent call last)
/tmp/ipykernel_43664/910117771.py in
3 x1 = torch.randn(50)
4 lazy_covar_matrix = covar_module(x1) # Returns a RootLazyTensor
----> 5 lazy_covar_matrix.diag()

~/anaconda3/lib/python3.8/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
57 kwargs_pkl = pickle.dumps(kwargs)
58 if not _is_in_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl):
---> 59 return _add_to_cache(self, cache_name, method(self, *args, **kwargs), *args, kwargs_pkl=kwargs_pkl)
60 return _get_from_cache(self, cache_name, *args, kwargs_pkl=kwargs_pkl)
61

~/anaconda3/lib/python3.8/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in diag(self)
309 expected_shape = self.shape[:-1]
310 if res.shape != expected_shape:
--> 311 raise RuntimeError(
312 "The kernel {} is not equipped to handle and diag. Expected size {}. "
313 "Got size {}".format(self.kernel.class.name, expected_shape, res.shape)

RuntimeError: The kernel PeriodicKernel is not equipped to handle and diag. Expected size torch.Size([50]). Got size torch.Size([50, 50])

Expected Behavior

Should return the diagonal of the kernel.

System information

Please complete the following information:
GPyTorch Version (run print(gpytorch.__version__)
1.6.0
PyTorch Version (run print(torch.__version__)
1.10.0
Computer OS
Windows 10

[gpleiss]

Thanks for catching this. I'm taking a look at this.