Repository: cornellius-gp/gpytorch
Pull request #631 (2019-04-09 17:44:12)

RBFKernelGrad fixes and improvements
- RBFKernelGrad now uses JIT for kernel postprocessing
- RBFKernelGrad now supports multibatch
- Use base kernel test cases for RBFKernelGrad
- Fix an issue where RBFKernelGrad was modifying `self.lengthscale` in-place.


Patch:
diff --git a/gpytorch/kernels/cosine_kernel.py b/gpytorch/kernels/cosine_kernel.py
index e6f88ec5..106363a1 100644
--- a/gpytorch/kernels/cosine_kernel.py
+++ b/gpytorch/kernels/cosine_kernel.py
@@ -82,7 +82,7 @@ class CosineKernel(Kernel):
 
     def _set_period_length(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_period_length)
+            value = torch.tensor(value)
 
         self.initialize(raw_period_length=self.raw_period_length_constraint.inverse_transform(value))
 
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index d6148302..d6fa7d58 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -272,7 +272,7 @@ class Kernel(Module):
             raise RuntimeError("Kernel has no lengthscale.")
 
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_lengthscale)
+            value = torch.tensor(value)
 
         self.initialize(raw_lengthscale=self.raw_lengthscale_constraint.inverse_transform(value))
 
diff --git a/gpytorch/kernels/linear_kernel.py b/gpytorch/kernels/linear_kernel.py
index e4a8dc8c..21bc324d 100644
--- a/gpytorch/kernels/linear_kernel.py
+++ b/gpytorch/kernels/linear_kernel.py
@@ -91,7 +91,7 @@ class LinearKernel(Kernel):
 
     def _set_variance(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_variance)
+            value = torch.tensor(value)
         self.initialize(raw_variance=self.raw_variance_constraint.inverse_transform(value))
 
     def forward(self, x1, x2, diag=False, batch_dims=None, **params):
diff --git a/gpytorch/kernels/periodic_kernel.py b/gpytorch/kernels/periodic_kernel.py
index a228b6f8..60e4385e 100644
--- a/gpytorch/kernels/periodic_kernel.py
+++ b/gpytorch/kernels/periodic_kernel.py
@@ -99,7 +99,7 @@ class PeriodicKernel(Kernel):
 
     def _set_period_length(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_period_length)
+            value = torch.tensor(value)
         self.initialize(raw_period_length=self.raw_period_length_constraint.inverse_transform(value))
 
     def forward(self, x1, x2, **params):
diff --git a/gpytorch/kernels/rbf_kernel_grad.py b/gpytorch/kernels/rbf_kernel_grad.py
index 700a96b4..f8b3599c 100644
--- a/gpytorch/kernels/rbf_kernel_grad.py
+++ b/gpytorch/kernels/rbf_kernel_grad.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from .rbf_kernel import RBFKernel
+from .rbf_kernel import RBFKernel, postprocess_rbf
 import torch
 from ..lazy.kronecker_product_lazy_tensor import KroneckerProductLazyTensor
 
@@ -51,6 +51,7 @@ class RBFKernelGrad(RBFKernel):
     """
     def forward(self, x1, x2, diag=False, **params):
         batch_shape = x1.shape[:-2]
+        n_batch_dims = len(batch_shape)
         n1, d = x1.shape[-2:]
         n2 = x2.shape[-2]
 
@@ -63,31 +64,37 @@ class RBFKernelGrad(RBFKernel):
 
             # Form all possible rank-1 products for the gradient and Hessian blocks
             outer = x1_.view(*batch_shape, n1, 1, d) - x2_.view(*batch_shape, 1, n2, d)
-            outer = outer / self.lengthscale
+            outer = outer / self.lengthscale.unsqueeze(-2)
             outer = torch.transpose(outer, -1, -2).contiguous()
 
             # 1) Kernel block
-            diff = self.covar_dist(x1_, x2_, square_dist=True, **params)
-            K_11 = diff.div_(-2).exp_()
+            diff = self.covar_dist(
+                x1_,
+                x2_,
+                square_dist=True,
+                dist_postprocess_func=postprocess_rbf,
+                **params
+            )
+            K_11 = diff
             K[..., :n1, :n2] = K_11
 
             # 2) First gradient block
             outer1 = outer.view(*batch_shape, n1, n2 * d)
-            K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])
+            K[..., :n1, n2:] = outer1 * K_11.repeat([*([1] * (n_batch_dims + 1)), d])
 
             # 3) Second gradient block
             outer2 = outer.transpose(-1, -3).contiguous().view(*batch_shape, n2, n1 * d)
             outer2 = outer2.transpose(-1, -2)
-            K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])
+            K[..., n1:, :n2] = -outer2 * K_11.repeat([*([1] * n_batch_dims), d, 1])
 
             # 4) Hessian block
-            outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])
+            outer3 = outer1.repeat([*([1] * n_batch_dims), d, 1]) * outer2.repeat([*([1] * (n_batch_dims + 1)), d])
             kp = KroneckerProductLazyTensor(
-                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow_(2),
+                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1) / self.lengthscale.pow(2),
                 torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat(*batch_shape, 1, 1)
             )
             chain_rule = kp.evaluate() - outer3
-            K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])
+            K[..., n1:, n2:] = chain_rule * K_11.repeat([*([1] * n_batch_dims), d, d])
 
             # Symmetrize for stability
             if n1 == n2 and torch.eq(x1, x2).all():
@@ -117,7 +124,4 @@ class RBFKernelGrad(RBFKernel):
         `d` dimensions, RBFKernelGrad returns an `n_1(d+1) x n_2(d+1)` kernel matrix.
         """
         non_batch_size = ((x1.size(-1) + 1) * x1.size(-2), (x2.size(-1) + 1) * x2.size(-2))
-        if x1.ndimension() == 3:
-            return torch.Size((x1.size(0),) + non_batch_size)
-        else:
-            return torch.Size(non_batch_size)
+        return torch.Size((*x1.shape[:-2], *non_batch_size))
diff --git a/gpytorch/kernels/scale_kernel.py b/gpytorch/kernels/scale_kernel.py
index 65f04510..e235495b 100644
--- a/gpytorch/kernels/scale_kernel.py
+++ b/gpytorch/kernels/scale_kernel.py
@@ -75,7 +75,7 @@ class ScaleKernel(Kernel):
 
     def _set_outputscale(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_outputscale)
+            value = torch.tensor(value)
         self.initialize(raw_outputscale=self.raw_outputscale_constraint.inverse_transform(value))
 
     def forward(self, x1, x2, batch_dims=None, diag=False, **params):
diff --git a/gpytorch/kernels/spectral_mixture_kernel.py b/gpytorch/kernels/spectral_mixture_kernel.py
index fd8cb55c..ad63602f 100644
--- a/gpytorch/kernels/spectral_mixture_kernel.py
+++ b/gpytorch/kernels/spectral_mixture_kernel.py
@@ -115,7 +115,7 @@ class SpectralMixtureKernel(Kernel):
 
     def _set_mixture_scales(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_mixture_scales)
+            value = torch.tensor(value)
         self.initialize(raw_mixture_scales=self.raw_mixture_scales_constraint.inverse_transform(value))
 
     @property
@@ -128,7 +128,7 @@ class SpectralMixtureKernel(Kernel):
 
     def _set_mixture_means(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_mixture_means)
+            value = torch.tensor(value)
         self.initialize(raw_mixture_means=self.raw_mixture_means_constraint.inverse_transform(value))
 
     @property
@@ -141,7 +141,7 @@ class SpectralMixtureKernel(Kernel):
 
     def _set_mixture_weights(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_mixture_weights)
+            value = torch.tensor(value)
         self.initialize(raw_mixture_weights=self.raw_mixture_weights_constraint.inverse_transform(value))
 
     def initialize_from_data(self, train_x, train_y, **kwargs):
diff --git a/gpytorch/likelihoods/multitask_gaussian_likelihood.py b/gpytorch/likelihoods/multitask_gaussian_likelihood.py
index b96713e9..e876ec2a 100644
--- a/gpytorch/likelihoods/multitask_gaussian_likelihood.py
+++ b/gpytorch/likelihoods/multitask_gaussian_likelihood.py
@@ -183,7 +183,7 @@ class MultitaskGaussianLikelihood(_MultitaskGaussianLikelihoodBase):
 
     def _set_noise(self, value):
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_noise)
+            value = torch.tensor(value)
         self.initialize(raw_noise=self.raw_noise_constraint.inverse_transform(value))
 
     def _shaped_noise_covar(self, base_shape, *params):
diff --git a/gpytorch/likelihoods/noise_models.py b/gpytorch/likelihoods/noise_models.py
index da7061e7..a198e8b1 100644
--- a/gpytorch/likelihoods/noise_models.py
+++ b/gpytorch/likelihoods/noise_models.py
@@ -39,7 +39,7 @@ class _HomoskedasticNoiseBase(Noise):
 
     def _set_noise(self, value: Tensor) -> None:
         if not torch.is_tensor(value):
-            value = torch.as_tensor(value).to(self.raw_noise)
+            value = torch.tensor(value)
         self.initialize(raw_noise=self.raw_noise_constraint.inverse_transform(value))
 
     def forward(self, *params: Any, shape: Optional[torch.Size] = None) -> DiagLazyTensor:


Test patch:
diff --git a/test/kernels/test_rbf_kernel_grad.py b/test/kernels/test_rbf_kernel_grad.py
index 486e33fa..258be919 100644
--- a/test/kernels/test_rbf_kernel_grad.py
+++ b/test/kernels/test_rbf_kernel_grad.py
@@ -3,9 +3,16 @@
 import torch
 import unittest
 from gpytorch.kernels import RBFKernelGrad
+from test.kernels._base_kernel_test_case import BaseKernelTestCase
 
 
-class TestRBFKernelGrad(unittest.TestCase):
+class TestRBFKernelGrad(unittest.TestCase, BaseKernelTestCase):
+    def create_kernel_no_ard(self, **kwargs):
+        return RBFKernelGrad(**kwargs)
+
+    def create_kernel_ard(self, num_dims, **kwargs):
+        return RBFKernelGrad(ard_num_dims=num_dims, **kwargs)
+
     def test_kernel(self, cuda=False):
         a = torch.tensor([[[1, 2], [2, 4]]], dtype=torch.float)
         b = torch.tensor([[[1, 3], [0, 4]]], dtype=torch.float)