Implement diag mode for PolynomialKernelGrad
@@ -25,8 +25,16 @@ class PolynomialKernelGrad(PolynomialKernel):
n2 = x2.shape[-2]
if diag:
- raise RuntimeError("None done yet")
- K11 = ((x1 * x2).sum(dim=-1) + self.offset).pow(self.power)
+ base_diag = ((x1 * x2).sum(dim=-1) + self.offset)
+ K11_diag = base_diag.pow(self.power)
+
+ all_outers_diag = (x1 * x2).transpose(-2, -1).contiguous().view(*batch_shape, -1)
+ 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)
+
+ 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)
