Repository: cornellius-gp/gpytorch

Issue:
ARD for RBFKernelGrad
Hi,
Is it possible to use ARD when training a GP model with derivative data (RBFkernelgrad)? Would it be possible to add support for that?


Discussion:
@Ushnish-Sengupta David has agreed to add ARD to the `RBFGradKernel`!
Oh amazing thanks, looking forward to that!

Resolving patch:
diff --git a/gpytorch/kernels/rbf_kernel_grad.py b/gpytorch/kernels/rbf_kernel_grad.py
--- a/gpytorch/kernels/rbf_kernel_grad.py
+++ b/gpytorch/kernels/rbf_kernel_grad.py
@@ -62,15 +62,15 @@ def forward(self, x1, x2, diag=False, **params):
             _, n2, _ = x2.size()
 
         K = torch.zeros(b, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)  # batch x n1(d+1) x n2(d+1)
-        ell = self.lengthscale.squeeze(-1)
 
         if not diag:
             # Scale the inputs by the lengthscale (for stability)
-            x1_ = x1 / ell
-            x2_ = x2 / ell
+            x1_ = x1.div(self.lengthscale)
+            x2_ = x2.div(self.lengthscale)
 
             # Form all possible rank-1 products for the gradient and Hessian blocks
             outer = x1_.view([b, n1, 1, d]) - x2_.view([b, 1, n2, d])
+            outer = outer / self.lengthscale
             outer = torch.transpose(outer, -1, -2).contiguous()
 
             # 1) Kernel block
@@ -79,21 +79,21 @@ def forward(self, x1, x2, diag=False, **params):
             K[..., :n1, :n2] = K_11
 
             # 2) First gradient block
-            outer1 = outer.view([b, n1, n2 * d]) / ell
+            outer1 = outer.view([b, n1, n2 * d])
             K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])
 
             # 3) Second gradient block
             outer2 = outer.transpose(-1, -3).contiguous().view([b, n2, n1 * d])
-            outer2 = outer2.transpose(-1, -2) / ell
+            outer2 = outer2.transpose(-1, -2)
             K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])
 
             # 4) Hessian block
             outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])
             kp = KroneckerProductLazyTensor(
-                torch.eye(d, d, device=x1.device, dtype=x1.dtype),
-                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype)
+                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1]) / self.lengthscale.pow_(2),
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1])
             )
-            chain_rule = kp.evaluate() / ell.pow(2) - outer3
+            chain_rule = kp.evaluate() - outer3
             K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])
 
             # Symmetrize for stability
@@ -107,12 +107,13 @@ def forward(self, x1, x2, diag=False, **params):
 
             return K
 
-        else:  # TODO: This will change when ARD is supported
+        else:
             if not (n1 == n2 and torch.eq(x1, x2).all()):
                 raise RuntimeError("diag=True only works when x1 == x2")
 
             kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)
-            grad_diag = torch.ones(1, n2 * d, device=x1.device, dtype=x1.dtype) / (ell.pow(2))
+            grad_diag = torch.ones(b, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow_(2)
+            grad_diag = grad_diag.transpose(-1, -2).contiguous().view([b, n2 * d])
             k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)
             pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))
             return k_diag[..., pi]


Test patch: