[Bug] _diagonal() fails when covariance is computed using keops kernels

[ancorso]

🐛 Bug

Calls to _diagnonal() error when keops kernels are used. This showed up when tring to use keops kernels in a Multitask/Multi-output approximate GP with SVI. But below is a MWE of the failure.

To reproduce

import gpytorch
import torch

train_x = torch.rand(100,2)
cov1 = gpytorch.kernels.RBFKernel()(train_x)
cov2 = gpytorch.kernels.keops.RBFKernel()(train_x)

cov1._diagonal() # Returns a tensor of shape (100,)
cov2._diagonal() # RuntimeError: The kernel RBFKernel is not equipped to handle and diag. Expected size torch.Size([100]). Got size torch.Size([100, 100])

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/anthonycorso/Workspace/gpytorch/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 "/Users/anthonycorso/Workspace/gpytorch/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 25, in wrapped
    output = method(self, *args, **kwargs)
  File "/Users/anthonycorso/Workspace/gpytorch/gpytorch/lazy/lazy_evaluated_kernel_tensor.py", line 126, in _diagonal
    raise RuntimeError(
RuntimeError: The kernel RBFKernel is not equipped to handle and diag. Expected size torch.Size([100]). Got size torch.Size([100, 100])

Expected Behavior

The output of cov2_diagonal() should match the output of cov1._diagonal(), as there should be no difference between built-in and keops kernels

System information

• GPyTorch Version = v1.12
• torch version = 2.4.0
• MacOS

Additional context

The problem can be fixed by handling the diag keyword argument in the forward call of the keops kernels. For example

    def forward(self, x1, x2, diag=False, **kwargs):
        x1_ = x1 / self.lengthscale
        x2_ = x2 / self.lengthscale
        # return KernelLinearOperator inst only when calculating the whole covariance matrix
        K = KernelLinearOperator(x1_, x2_, covar_func=_covar_func, **kwargs)    
        return K.diagonal() if diag else K   

I am happy to submit a PR with this change (and a similar change for the other keops kernels, but I'm not 100% sure this is the "right" solution or if there is a more elegant intervention.

[gpleiss]

I opened a PR with a hotfix

[ancorso]

thank you! will close when PR is merged!