Refactor to Arguments to more closely match torch.tensor (#1746)

* refactor to arguments

* switch to args, kwargs constructor

* add _to_helper to broadcasting util

Author: wjmaddox

gpytorch/lazy/cat_lazy_tensor.py

@@ -3,7 +3,7 @@
 import torch
 
 from .. import settings
-from ..utils.broadcasting import _matmul_broadcast_shape, _mul_broadcast_shape
+from ..utils.broadcasting import _matmul_broadcast_shape, _mul_broadcast_shape, _to_helper
 from ..utils.deprecation import bool_compat
 from ..utils.getitem import _noop_index
 from .lazy_tensor import LazyTensor, delazify
@@ -364,15 +364,22 @@ def devices(self):
     def device_count(self):
         return len(set(self.devices))
 
-    def to(self, device_id):
+    def to(self, *args, **kwargs):
         """
-        returns a new CatLazyTensor with device_id as the output_device
+        Returns a new CatLazyTensor with device as the output_device and dtype
+        as the dtype.
         Warning: this does not move the LazyTensors in this CatLazyTensor to
-        device_id
+        device.
         """
-        new_kwargs = dict(self._kwargs)
-        new_kwargs["output_device"] = device_id
-        return self.__class__(*self._args, **new_kwargs)
+        device, dtype = _to_helper(*args, **kwargs)
+
+        new_kwargs = {**self._kwargs, "output_device": device}
+        res = self.__class__(*self._args, **new_kwargs)
+
+        if dtype is not None:
+            res = res.type(dtype)
+
+        return res
 
     def all_to(self, device_id):
         """

gpytorch/lazy/lazy_tensor.py

@@ -19,7 +19,7 @@
 from ..functions._matmul import Matmul
 from ..functions._root_decomposition import RootDecomposition
 from ..functions._sqrt_inv_matmul import SqrtInvMatmul
-from ..utils.broadcasting import _matmul_broadcast_shape, _mul_broadcast_shape
+from ..utils.broadcasting import _matmul_broadcast_shape, _mul_broadcast_shape, _to_helper
 from ..utils.cholesky import psd_safe_cholesky
 from ..utils.deprecation import _deprecate_renamed_methods
 from ..utils.errors import CachingError
@@ -1859,25 +1859,31 @@ def symeig(self, eigenvectors: bool = False) -> Tuple[Tensor, Optional["LazyTens
             pass
         return self._symeig(eigenvectors=eigenvectors)
 
-    def to(self, device_id):
+    def to(self, *args, **kwargs):
         """
-        A device-agnostic method of moving the lazy_tensor to the specified device.
+        A device-agnostic method of moving the lazy_tensor to the specified device or dtype.
+        Note that we do NOT support non_blocking or other `torch.to` options other than
+        device and dtype and these options will be silently ignored.
 
         Args:
-            device_id (:obj: `torch.device`): Which device to use (GPU or CPU).
+            device (:obj: `torch.device`): Which device to use (GPU or CPU).
+            dtype (:obj: `torch.dtype`): Which dtype to use (double, float, or half).
         Returns:
             :obj:`~gpytorch.lazy.LazyTensor`: New LazyTensor identical to self on specified device
         """
+
+        device, dtype = _to_helper(*args, **kwargs)
+
         new_args = []
         new_kwargs = {}
         for arg in self._args:
             if hasattr(arg, "to"):
-                new_args.append(arg.to(device_id))
+                new_args.append(arg.to(dtype=dtype, device=device))
             else:
                 new_args.append(arg)
         for name, val in self._kwargs.items():
             if hasattr(val, "to"):
-                new_kwargs[name] = val.to(device_id)
+                new_kwargs[name] = val.to(dtype=dtype, device=device)
             else:
                 new_kwargs[name] = val
         return self.__class__(*new_args, **new_kwargs)

gpytorch/utils/broadcasting.py

@@ -69,3 +69,27 @@ def _pad_with_singletons(obj, num_singletons_before=0, num_singletons_after=0):
     """
     new_shape = [1] * num_singletons_before + list(obj.shape) + [1] * num_singletons_after
     return obj.view(*new_shape)
+
+
+def _to_helper(*args, **kwargs):
+    """
+    Silently plucks out dtype and devices from a list.
+
+    Example:
+        >>> dtype, device = _to_helper(dtype=torch.float, device=torch.device("cpu"))
+        >>> dtype, device = _to_helper(torch.float, torch.device("cpu"))
+    """
+    dtype = kwargs.pop("dtype", None)
+    device = kwargs.pop("device", None)
+
+    if dtype is None:
+        dtype_list = [x for x in args if type(x) is torch.dtype]
+        if len(dtype_list) > 0:
+            dtype = dtype_list[0]
+
+    if device is None:
+        device_list = [x for x in args if type(x) is torch.device]
+        if len(device_list) > 0:
+            device = device_list[0]
+
+    return device, dtype