Merge pull request #1726 from wjmaddox/half_float_calls

Add half and float calls to lazy tensor

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/lazy/lazy_tensor.py

@@ -3,6 +3,7 @@
 import math
 import warnings
 from abc import ABC, abstractmethod
+from copy import deepcopy
 from typing import Optional, Tuple
 
 import torch
@@ -30,6 +31,8 @@
 from ..utils.warnings import NumericalWarning
 from .lazy_tensor_representation_tree import LazyTensorRepresentationTree
 
+_TYPES_DICT = {torch.float: "float", torch.half: "half", torch.double: "double"}
+
 
 class LazyTensor(ABC):
     r"""
@@ -1063,19 +1066,7 @@ def double(self, device_id=None):
         """
         This method operates identically to :func:`torch.Tensor.double`.
         """
-        new_args = []
-        new_kwargs = {}
-        for arg in self._args:
-            if hasattr(arg, "double"):
-                new_args.append(arg.double())
-            else:
-                new_args.append(arg)
-        for name, val in self._kwargs.items():
-            if hasattr(val, "double"):
-                new_kwargs[name] = val.double()
-            else:
-                new_kwargs[name] = val
-        return self.__class__(*new_args, **new_kwargs)
+        return self.type(torch.double)
 
     @property
     def dtype(self):
@@ -1122,6 +1113,18 @@ def evaluate_kernel(self):
         """
         return self.representation_tree()(*self.representation())
 
+    def float(self, device_id=None):
+        """
+        This method operates identically to :func:`torch.Tensor.float`.
+        """
+        return self.type(torch.float)
+
+    def half(self, device_id=None):
+        """
+        This method operates identically to :func:`torch.Tensor.half`.
+        """
+        return self.type(torch.half)
+
     def inv_matmul(self, right_tensor, left_tensor=None):
         r"""
         Computes a linear solve (w.r.t self = :math:`A`) with several right hand sides :math:`R`.
@@ -1924,6 +1927,32 @@ def transpose(self, dim1, dim2):
 
         return res
 
+    def type(self, dtype):
+        """
+        This method operates similarly to :func:`torch.Tensor.type`.
+        """
+        attr_flag = _TYPES_DICT[dtype]
+
+        new_args = []
+        new_kwargs = {}
+        for arg in self._args:
+            if hasattr(arg, attr_flag):
+                try:
+                    new_args.append(arg.clone().to(dtype))
+                except AttributeError:
+                    new_args.append(deepcopy(arg).to(dtype))
+            else:
+                new_args.append(arg)
+        for name, val in self._kwargs.items():
+            if hasattr(val, attr_flag):
+                try:
+                    new_kwargs[name] = val.clone().to(dtype)
+                except AttributeError:
+                    new_kwargs[name] = deepcopy(val).to(dtype)
+            else:
+                new_kwargs[name] = val
+        return self.__class__(*new_args, **new_kwargs)
+
     def unsqueeze(self, dim):
         positive_dim = (self.dim() + dim + 1) if dim < 0 else dim
         if positive_dim > len(self.batch_shape):

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)