use type method

Author: wjmaddox

gpytorch/lazy/lazy_tensor.py

@@ -31,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"""
@@ -1064,25 +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"):
-                try:
-                    new_args.append(arg.clone().double())
-                except AttributeError:
-                    new_args.append(deepcopy(arg).double())
-            else:
-                new_args.append(arg)
-        for name, val in self._kwargs.items():
-            if hasattr(val, "double"):
-                try:
-                    new_kwargs[name] = val.clone().double()
-                except AttributeError:
-                    new_kwargs[name] = deepcopy(val).double()
-            else:
-                new_kwargs[name] = val
-        return self.__class__(*new_args, **new_kwargs)
+        return self.type(torch.double)
 
     @property
     def dtype(self):
@@ -1133,49 +1117,13 @@ def float(self, device_id=None):
         """
         This method operates identically to :func:`torch.Tensor.float`.
         """
-        new_args = []
-        new_kwargs = {}
-        for arg in self._args:
-            if hasattr(arg, "float"):
-                try:
-                    new_args.append(arg.clone().float())
-                except AttributeError:
-                    new_args.append(deepcopy(arg).float())
-            else:
-                new_args.append(arg)
-        for name, val in self._kwargs.items():
-            if hasattr(val, "float"):
-                try:
-                    new_kwargs[name] = val.clone().float()
-                except AttributeError:
-                    new_kwargs[name] = deepcopy(val).float()
-            else:
-                new_kwargs[name] = val
-        return self.__class__(*new_args, **new_kwargs)
+        return self.type(torch.float)
 
     def half(self, device_id=None):
         """
         This method operates identically to :func:`torch.Tensor.half`.
         """
-        new_args = []
-        new_kwargs = {}
-        for arg in self._args:
-            if hasattr(arg, "half"):
-                try:
-                    new_args.append(arg.clone().half())
-                except AttributeError:
-                    new_args.append(deepcopy(arg).half())
-            else:
-                new_args.append(arg)
-        for name, val in self._kwargs.items():
-            if hasattr(val, "half"):
-                try:
-                    new_kwargs[name] = val.clone().half()
-                except AttributeError:
-                    new_kwargs[name] = deepcopy(val).half()
-            else:
-                new_kwargs[name] = val
-        return self.__class__(*new_args, **new_kwargs)
+        return self.type(torch.half)
 
     def inv_matmul(self, right_tensor, left_tensor=None):
         r"""
@@ -1980,14 +1928,30 @@ def transpose(self, dim1, dim2):
         return res
 
     def type(self, dtype):
-        if dtype == torch.float:
-            return self.float()
-        elif dtype == torch.double:
-            return self.double()
-        elif dtype == torch.half:
-            return self.half()
-        else:
-            raise RuntimeError("Dtype", dtype, "not found.")
+        """
+        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