def METHOD_NAME(self):
    a = torch.tensor([[4, 1], [2, 2], [8, 0]], dtype=torch.float)
    b = torch.tensor([[0, 0], [2, 1], [1, 0]], dtype=torch.float)
    kernel = PolynomialKernel(power=3)
    kernel.eval()
    actual = torch.zeros(3, 3)
    for i in range(3):
        for j in range(3):
            actual[i, j] = (a[i].matmul(b[j]) + kernel.offset).pow(kernel.power)
    res = kernel(a, b).to_dense()
    self.assertLess(torch.norm(res - actual), 1e-5)
    # diag
    res = kernel(a, b).diagonal(dim1=-1, dim2=-2)
    actual = actual.diagonal(dim1=-1, dim2=-2)
    self.assertLess(torch.norm(res - actual), 1e-5)
    # batch_dims
    actual = torch.zeros(2, 3, 3)
    for i in range(2):
        actual[i] = kernel(a[:, i].unsqueeze(-1), b[:, i].unsqueeze(-1)).to_dense()
    res = kernel(a, b, last_dim_is_batch=True).to_dense()
    self.assertLess(torch.norm(res - actual), 1e-5)
    # batch_dims + diag
    res = kernel(a, b, last_dim_is_batch=True).diagonal(dim1=-1, dim2=-2)
    actual = torch.cat([actual[i].diagonal(dim1=-1, dim2=-2).unsqueeze(0) for i in range(actual.size(0))])
    self.assertLess(torch.norm(res - actual), 1e-5)