Fix diag calls to ScaleKernel
@@ -45,18 +45,24 @@ class ScaleKernel(Kernel):
>>> scaled_covar_module = gpytorch.kernels.ScaleKernel(base_covar_module)
>>> covar = scaled_covar_module(x) # Output: LazyTensor of size (10 x 10)
"""
-
def __init__(self, base_kernel, batch_size=1, log_outputscale_prior=None):
super(ScaleKernel, self).__init__(has_lengthscale=False, batch_size=batch_size)
self.base_kernel = base_kernel
self.register_parameter(
- name="log_outputscale", parameter=torch.nn.Parameter(torch.zeros(batch_size)), prior=log_outputscale_prior
+ name="log_outputscale",
+ parameter=torch.nn.Parameter(torch.zeros(batch_size)),
+ prior=log_outputscale_prior
)
@property
def outputscale(self):
return self.log_outputscale.exp()
+ def forward_diag(self, x1, x2):
+ outputscales = self.log_outputscale.exp()
+ orig_output = self.base_kernel.forward_diag(x1, x2)
+ return orig_output.squeeze(-1).mul(outputscales.unsqueeze(-1))
+
def forward(self, x1, x2):
outputscales = self.log_outputscale.exp()
orig_output = self.base_kernel(x1, x2)
