Repository: cornellius-gp/gpytorch
Pull request #225 (2018-08-08 17:07:31)

Add suite of unit tests for WhiteNoiseKernel and ZeroLazyVariable
This also fixes some issues that turned up with them in writing the unit tests for the batch setting. 

This PR begins work on #224. I think it is best to make PRs for unit tests as I make them rather than trying to do a big push for them all in one PR, since I may also be able to fix minor issues as I make them as in this PR.

cc/ @gpleiss @Balandat 


Patch:
diff --git a/gpytorch/kernels/white_noise_kernel.py b/gpytorch/kernels/white_noise_kernel.py
index 5cc1e040..0ebc7348 100644
--- a/gpytorch/kernels/white_noise_kernel.py
+++ b/gpytorch/kernels/white_noise_kernel.py
@@ -13,10 +13,25 @@ class WhiteNoiseKernel(Kernel):
         super(WhiteNoiseKernel, self).__init__()
         self.register_buffer("variances", variances)
 
+        # Enforce that variances are batch_size x data_size x task_size
+        # where batch_size and task_size are both 1 in the simplest case.
+        if variances.ndimension() == 3:
+            self.is_batch = True
+            self.variances = variances
+        elif variances.ndimension() == 2:
+            # By default, assume 2 dimensional data means data_size x task_size, and add a batch_size of 1
+            self.is_batch = False
+            self.variances = variances.unsqueeze(0)
+        else:
+            self.is_batch = False
+            self.variances = variances.unsqueeze(0).unsqueeze(-1)
+
     def forward(self, x1, x2):
-        if self.training:
-            return DiagLazyVariable(self.variances.view(-1).unsqueeze(0))
-        elif x1.size(-2) == x2.size(-2) and x1.size(-2) == self.variances.size(0) and torch.equal(x1, x2):
-            return DiagLazyVariable(self.variances.view(-1).unsqueeze(0))
+        if self.training and torch.equal(x1, x2):
+                # Reshape into a batch of batch_size diagonal matrices, each of which is
+                # (data_size * task_size) x (data_size * task_size)
+                return DiagLazyVariable(self.variances.view(self.variances.size(0), -1))
+        elif x1.size(-2) == x2.size(-2) and x1.size(-2) == self.variances.size(1) and torch.equal(x1, x2):
+                return DiagLazyVariable(self.variances.view(self.variances.size(0), -1))
         else:
             return ZeroLazyVariable(x1.size(-3), x1.size(-2), x2.size(-2))
diff --git a/gpytorch/lazy/zero_lazy_variable.py b/gpytorch/lazy/zero_lazy_variable.py
index fcb3ec9a..28b9fc31 100644
--- a/gpytorch/lazy/zero_lazy_variable.py
+++ b/gpytorch/lazy/zero_lazy_variable.py
@@ -124,3 +124,18 @@ class ZeroLazyVariable(LazyVariable):
 
     def __mul__(self, other):
         return self
+
+    def __getitem__(self, index):
+        index = list(index) if isinstance(index, tuple) else [index]
+        ndimension = self.ndimension()
+        index += [slice(None, None, None)] * (ndimension - len(index))
+        new_sizes = []
+        for ix, sub_index in enumerate(index):
+            if isinstance(sub_index, int):
+                continue
+            elif isinstance(sub_index, slice):
+                new_sizes += [len(torch.arange(self.size(ix))[sub_index])]
+            else:
+                new_sizes += [len(sub_index)]
+
+        return ZeroLazyVariable(*new_sizes)


Test patch:
diff --git a/test/kernels/test_white_noise_kernel.py b/test/kernels/test_white_noise_kernel.py
new file mode 100644
index 00000000..dd276f19
--- /dev/null
+++ b/test/kernels/test_white_noise_kernel.py
@@ -0,0 +1,78 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import torch
+import unittest
+from gpytorch.kernels import WhiteNoiseKernel
+
+
+class TestWhiteNoiseKernel(unittest.TestCase):
+    def test_computes_diag_train(self):
+        a = torch.Tensor([4, 2, 8]).view(3, 1)
+        variances = torch.randn(3)
+        kernel = WhiteNoiseKernel(variances=variances)
+        actual = torch.diag(variances)
+        res = kernel(a).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_computes_diag_train_batch(self):
+        a = torch.Tensor([[4, 2, 8], [4, 2, 8]]).view(2, 3, 1)
+        variances = torch.randn(2, 3, 1)
+        kernel = WhiteNoiseKernel(variances=variances)
+        actual = torch.cat(
+            (torch.diag(variances[0].squeeze(-1)).unsqueeze(0), torch.diag(variances[1].squeeze(-1)).unsqueeze(0))
+        )
+        res = kernel(a).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_computes_zero_eval(self):
+        a = torch.Tensor([4, 2, 8]).view(3, 1)
+        b = torch.Tensor([3, 7]).view(2, 1)
+        variances = torch.randn(3)
+        kernel = WhiteNoiseKernel(variances=variances)
+        kernel.eval()
+        actual_one = torch.zeros(3, 2)
+        actual_two = torch.zeros(2, 3)
+        res_one = kernel(a, b).evaluate()
+        res_two = kernel(b, a).evaluate()
+        self.assertLess(torch.norm(res_one - actual_one), 1e-5)
+        self.assertLess(torch.norm(res_two - actual_two), 1e-5)
+
+    def test_computes_zero_eval_batch(self):
+        a = torch.Tensor([[4, 2, 8], [4, 2, 8]]).view(2, 3, 1)
+        b = torch.Tensor([[3, 7], [3, 7]]).view(2, 2, 1)
+        variances = torch.randn(2, 3, 1)
+        kernel = WhiteNoiseKernel(variances=variances)
+        kernel.eval()
+        actual_one = torch.zeros(3, 2)
+        actual_two = torch.zeros(2, 3)
+        res_one = kernel(a, b).evaluate()
+        res_two = kernel(b, a).evaluate()
+        self.assertLess(torch.norm(res_one - actual_one), 1e-5)
+        self.assertLess(torch.norm(res_two - actual_two), 1e-5)
+
+    def test_computes_diag_eval(self):
+        a = torch.Tensor([4, 2, 8]).view(3, 1)
+        variances = torch.randn(3)
+        kernel = WhiteNoiseKernel(variances=variances)
+        kernel.eval()
+        actual = torch.diag(variances)
+        res = kernel(a).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+    def test_computes_diag_eval_batch(self):
+        a = torch.Tensor([[4, 2, 8], [4, 2, 8]]).view(2, 3, 1)
+        variances = torch.randn(2, 3, 1)
+        kernel = WhiteNoiseKernel(variances=variances)
+        kernel.eval()
+        actual = torch.cat(
+            (torch.diag(variances[0].squeeze(-1)).unsqueeze(0), torch.diag(variances[1].squeeze(-1)).unsqueeze(0))
+        )
+        res = kernel(a).evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-5)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/lazy/test_zero_lazy_variable.py b/test/lazy/test_zero_lazy_variable.py
new file mode 100644
index 00000000..6e696b91
--- /dev/null
+++ b/test/lazy/test_zero_lazy_variable.py
@@ -0,0 +1,42 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import torch
+import unittest
+from gpytorch.lazy import ZeroLazyVariable
+
+
+class TestZeroLazyVariable(unittest.TestCase):
+    def test_evaluate(self):
+        lv = ZeroLazyVariable(5, 4, 3)
+        actual = torch.zeros(5, 4, 3)
+        res = lv.evaluate()
+        self.assertLess(torch.norm(res - actual), 1e-4)
+
+    def test_getitem(self):
+        lv = ZeroLazyVariable(5, 4, 3)
+
+        res_one = lv[0].evaluate()
+        res_two = lv[:, 1, :].evaluate()
+        res_three = lv[:, :, 2].evaluate()
+
+        self.assertLess(torch.norm(res_one - torch.zeros(4, 3)), 1e-4)
+        self.assertLess(torch.norm(res_two - torch.zeros(5, 3)), 1e-4)
+        self.assertLess(torch.norm(res_three - torch.zeros(5, 4)), 1e-4)
+
+    def test_getitem_complex(self):
+        lv = ZeroLazyVariable(5, 4, 3)
+
+        res_one = lv[[0, 1]].evaluate()
+        res_two = lv[:, [0, 1], :].evaluate()
+        res_three = lv[:, :, [0, 2]].evaluate()
+
+        self.assertLess(torch.norm(res_one - torch.zeros(2, 4, 3)), 1e-4)
+        self.assertLess(torch.norm(res_two - torch.zeros(5, 2, 3)), 1e-4)
+        self.assertLess(torch.norm(res_three - torch.zeros(5, 4, 2)), 1e-4)
+
+
+if __name__ == "__main__":
+    unittest.main()