add float, double, half unit tests to lt testcase

Author: wjmaddox

gpytorch/lazy/interpolated_lazy_tensor.py

@@ -386,6 +386,13 @@ def diag(self):
         else:
             return super(InterpolatedLazyTensor, self).diag()
 
+    def double(self, device_id=None):
+        # We need to ensure that the indices remain integers.
+        new_lt = super().double(device_id=device_id)
+        new_lt.left_interp_indices = new_lt.left_interp_indices.type(torch.int64)
+        new_lt.right_interp_indices = new_lt.right_interp_indices.type(torch.int64)
+        return new_lt
+
     def matmul(self, tensor):
         # We're using a custom matmul here, because it is significantly faster than
         # what we get from the function factory.

gpytorch/test/lazy_tensor_test_case.py

@@ -475,6 +475,14 @@ def test_cholesky(self):
             self.assertAllClose(res, actual, **self.tolerances["cholesky"])
             # TODO: Check gradients
 
+    def test_double(self):
+        lazy_tensor = self.create_lazy_tensor()
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+
+        res = lazy_tensor.double()
+        actual = evaluated.double()
+        self.assertEqual(res.dtype, actual.dtype)
+
     def test_diag(self):
         lazy_tensor = self.create_lazy_tensor()
         evaluated = self.evaluate_lazy_tensor(lazy_tensor)
@@ -484,6 +492,25 @@ def test_diag(self):
         actual = actual.view(*lazy_tensor.batch_shape, -1)
         self.assertAllClose(res, actual, **self.tolerances["diag"])
 
+    def test_float(self):
+        lazy_tensor = self.create_lazy_tensor().double()
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+
+        res = lazy_tensor.float()
+        actual = evaluated.float()
+        self.assertEqual(res.dtype, actual.dtype)
+
+    def _test_half(self, lazy_tensor):
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+
+        res = lazy_tensor.half()
+        actual = evaluated.half()
+        self.assertEqual(res.dtype, actual.dtype)
+
+    def test_half(self):
+        lazy_tensor = self.create_lazy_tensor()
+        self._test_half(lazy_tensor)
+
     def test_inv_matmul_vector(self, cholesky=False):
         lazy_tensor = self.create_lazy_tensor()
         rhs = torch.randn(lazy_tensor.size(-1))

test/lazy/test_interpolated_lazy_tensor.py

@@ -35,8 +35,8 @@ def create_lazy_tensor(self):
         )
 
     def evaluate_lazy_tensor(self, lazy_tensor):
-        left_matrix = torch.zeros(4, 6)
-        right_matrix = torch.zeros(4, 6)
+        left_matrix = torch.zeros(4, 6, dtype=lazy_tensor.dtype)
+        right_matrix = torch.zeros(4, 6, dtype=lazy_tensor.dtype)
         left_matrix.scatter_(1, lazy_tensor.left_interp_indices, lazy_tensor.left_interp_values)
         right_matrix.scatter_(1, lazy_tensor.right_interp_indices, lazy_tensor.right_interp_values)
 
@@ -75,8 +75,8 @@ def evaluate_lazy_tensor(self, lazy_tensor):
         left_matrix_comps = []
         right_matrix_comps = []
         for i in range(5):
-            left_matrix_comp = torch.zeros(4, 6)
-            right_matrix_comp = torch.zeros(4, 6)
+            left_matrix_comp = torch.zeros(4, 6, dtype=lazy_tensor.dtype)
+            right_matrix_comp = torch.zeros(4, 6, dtype=lazy_tensor.dtype)
             left_matrix_comp.scatter_(1, lazy_tensor.left_interp_indices[i], lazy_tensor.left_interp_values[i])
             right_matrix_comp.scatter_(1, lazy_tensor.right_interp_indices[i], lazy_tensor.right_interp_values[i])
             left_matrix_comps.append(left_matrix_comp.unsqueeze(0))
@@ -121,8 +121,8 @@ def evaluate_lazy_tensor(self, lazy_tensor):
         right_matrix_comps = []
         for i in range(2):
             for j in range(5):
-                left_matrix_comp = torch.zeros(4, 6)
-                right_matrix_comp = torch.zeros(4, 6)
+                left_matrix_comp = torch.zeros(4, 6, dtype=lazy_tensor.dtype)
+                right_matrix_comp = torch.zeros(4, 6, dtype=lazy_tensor.dtype)
                 left_matrix_comp.scatter_(
                     1, lazy_tensor.left_interp_indices[i, j], lazy_tensor.left_interp_values[i, j]
                 )

test/lazy/test_lazy_evaluated_kernel_tensor.py

@@ -127,6 +127,13 @@ def test_getitem_tensor_index(self):
     def test_quad_form_derivative(self):
         pass
 
+    def test_half(self):
+        # many transform operations aren't supported in half so we overwrite
+        # this test
+        lazy_tensor = self.create_lazy_tensor()
+        lazy_tensor.kernel.raw_lengthscale_constraint.transform = lambda x: x + 0.1
+        self._test_half(lazy_tensor)
+
 
 class TestLazyEvaluatedKernelTensorMultitaskBatch(TestLazyEvaluatedKernelTensorBatch):
     seed = 0
@@ -140,6 +147,13 @@ def create_lazy_tensor(self):
     def test_inv_matmul_matrix_with_checkpointing(self):
         pass
 
+    def test_half(self):
+        # many transform operations aren't supported in half so we overwrite
+        # this test
+        lazy_tensor = self.create_lazy_tensor()
+        lazy_tensor.kernel.data_covar_module.raw_lengthscale_constraint.transform = lambda x: x + 0.1
+        self._test_half(lazy_tensor)
+
 
 class TestLazyEvaluatedKernelTensorAdditive(TestLazyEvaluatedKernelTensorBatch):
     seed = 0
@@ -162,3 +176,10 @@ def evaluate_lazy_tensor(self, lazy_tensor):
 
     def test_inv_matmul_matrix_with_checkpointing(self):
         pass
+
+    def test_half(self):
+        # many transform operations aren't supported in half so we overwrite
+        # this test
+        lazy_tensor = self.create_lazy_tensor()
+        lazy_tensor.kernel.base_kernel.raw_lengthscale_constraint.transform = lambda x: x + 0.1
+        self._test_half(lazy_tensor)