Added th.device to represent both str and torch.device. (#357)

* Added th.device to represent both str and torch.device.

* _TheseusDevice now includes optional as well.

* Renamed TheseusDevice as DeviceType.

MANIFEST.in

@@ -1,3 +1,4 @@
 include LICENSE README.md
 include requirements/*.txt
+include py.typed
 graft theseus/extlib

theseus/__init__.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 __version__ = "0.1.2"
 
+from .constants import DeviceType as DeviceType
 
 from .core import (  # usort: skip
     AutoDiffCostFunction,

theseus/constants.py

@@ -2,7 +2,7 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-
+from typing import Optional, Union
 import math
 
 import torch
@@ -56,3 +56,5 @@ def __getitem__(self, dtype):
 _SE3_NEAR_ZERO_EPS = EPSDict(float32_eps=1e-2, float64_eps=5e-3)
 
 _SE3_HAT_EPS = EPSDict(float32_eps=5e-6, float64_eps=5e-7)
+
+DeviceType = Optional[Union[str, torch.device]]

theseus/core/objective.py

@@ -10,6 +10,7 @@
 
 import torch
 
+from theseus.constants import DeviceType
 from theseus.core.theseus_function import TheseusFunction
 from theseus.geometry.manifold import Manifold
 
@@ -55,7 +56,7 @@ def __init__(self, dtype: Optional[torch.dtype] = None):
 
         self._batch_size: Optional[int] = None
 
-        self.device: torch.device = torch.device("cpu")
+        self.device: DeviceType = torch.device("cpu")
 
         self.dtype: Optional[torch.dtype] = dtype or torch.get_default_dtype()
 

theseus/core/variable.py

@@ -7,6 +7,7 @@
 from typing import Optional, Sequence, Union
 
 import torch
+from theseus.constants import DeviceType
 
 
 class Variable:
@@ -107,12 +108,15 @@ def __setitem__(self, item, value):
 # In this case, the device, dtype and name can be specified.
 def as_variable(
     value: Union[float, Sequence[float], torch.Tensor, Variable],
-    device: Optional[torch.device] = None,
+    device: DeviceType = None,
     dtype: Optional[torch.dtype] = None,
     name: Optional[str] = None,
 ) -> Variable:
     if isinstance(value, Variable):
         return value
+    if isinstance(device, str):
+        # mypy doesn't like passing str to torch.as_tensor so convert
+        device = torch.device(device)
     tensor = torch.as_tensor(value, dtype=dtype, device=device)
     if isinstance(value, float):
         tensor = tensor.view(1, 1)

theseus/embodied/kinematics/kinematics_model.py

@@ -8,6 +8,7 @@
 
 import torch
 
+from theseus.constants import DeviceType
 from theseus.geometry import SE3, LieGroup, Point2, Vector
 
 RobotModelInput = Union[torch.Tensor, Vector]
@@ -36,7 +37,7 @@ def forward_kinematics(self, robot_pose: RobotModelInput) -> Dict[str, LieGroup]
 
 
 class UrdfRobotModel(KinematicsModel):
-    def __init__(self, urdf_path: str, device: Optional[str] = None):
+    def __init__(self, urdf_path: str, device: DeviceType = None):
         try:
             import differentiable_robot_model as drm
         except ModuleNotFoundError as e:

theseus/geometry/lie_group.py

@@ -8,6 +8,7 @@
 
 import torch
 
+from theseus.constants import DeviceType
 from theseus.geometry.manifold import Manifold
 
 
@@ -51,7 +52,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "LieGroup":
         pass
@@ -62,7 +63,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "LieGroup":
         pass

theseus/geometry/point_types.py

@@ -7,6 +7,7 @@
 
 import torch
 
+from theseus.constants import DeviceType
 from .vector import Vector
 
 
@@ -41,7 +42,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Point2":
         if len(size) != 1:
@@ -62,7 +63,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Point2":
         if len(size) != 1:
@@ -143,7 +144,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Point3":
         if len(size) != 1:
@@ -164,7 +165,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Point3":
         if len(size) != 1:

theseus/geometry/se2.py

@@ -44,7 +44,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SE2":
         if len(size) != 1:
@@ -66,7 +66,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SE2":
         if len(size) != 1:

theseus/geometry/se3.py

@@ -8,7 +8,6 @@
 
 import torch
 
-import theseus
 import theseus.constants
 
 from .lie_group import LieGroup
@@ -45,7 +44,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SE3":
         if len(size) != 1:
@@ -73,7 +72,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SE3":
         if len(size) != 1:

theseus/geometry/so2.py

@@ -43,7 +43,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SO2":
         if len(size) != 1:
@@ -67,7 +67,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SO2":
         if len(size) != 1:

theseus/geometry/so3.py

@@ -43,7 +43,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SO3":
         # Reference:
@@ -80,7 +80,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SO3":
         if len(size) != 1:

theseus/geometry/vector.py

@@ -7,6 +7,7 @@
 
 import torch
 
+from theseus.constants import DeviceType
 from theseus.geometry.manifold import Manifold
 
 from .lie_group import LieGroup
@@ -46,7 +47,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Vector":
         if len(size) != 2:
@@ -66,7 +67,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Vector":
         if len(size) != 2:

theseus/optimizer/linear/utils.py

@@ -4,6 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 from typing import Tuple, Union
 import torch
+from theseus.constants import DeviceType
 
 
 # See Nocedal and Wright, Numerical Optimization, pp. 260 and 261
@@ -13,7 +14,7 @@ def convert_to_alpha_beta_damping_tensors(
     damping_eps: float,
     ellipsoidal_damping: bool,
     batch_size: int,
-    device: torch.device,
+    device: DeviceType,
     dtype: torch.dtype,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     damping = torch.as_tensor(damping).to(device=device, dtype=dtype)

theseus/theseus_layer.py

@@ -18,7 +18,7 @@
     Variable,
     Vectorize,
 )
-from theseus.constants import __FROM_THESEUS_LAYER_TOKEN__
+from theseus.constants import __FROM_THESEUS_LAYER_TOKEN__, DeviceType
 from theseus.geometry import LieGroup, Manifold
 from theseus.optimizer import Optimizer, OptimizerInfo
 from theseus.optimizer.linear import LinearSolver
@@ -139,7 +139,7 @@ def to(self, *args, **kwargs):
         self.objective.to(*args, **kwargs)
 
     @property
-    def device(self) -> torch.device:
+    def device(self) -> DeviceType:
         return self.objective.device
 
     @property

theseus/utils/examples/motion_planning/motion_planner.py

@@ -231,7 +231,7 @@ def __init__(
         max_optim_iters: int,
         step_size: float = 1.0,
         objective: Optional[MotionPlannerObjective] = None,
-        device: str = "cpu",
+        device: th.DeviceType = "cpu",
         dtype: torch.dtype = torch.double,
         # The following are only used if objective is None
         map_size: Optional[int] = None,

theseus/utils/examples/pose_graph/dataset.py

@@ -182,7 +182,7 @@ def __init__(
         edges: List[PoseGraphEdge],
         gt_poses: Optional[Union[List[th.SE2], List[th.SE3]]] = None,
         batch_size: int = 1,
-        device: Optional[torch.device] = None,
+        device: th.DeviceType = None,
     ):
         dataset_sizes: List[int] = [pose.shape[0] for pose in poses]
         if gt_poses is not None:

theseus/utils/examples/tactile_pose_estimation/misc.py

@@ -5,6 +5,8 @@
 import numpy as np
 import torch
 
+import theseus as th
+
 # ----------------------------------------------------------------------------------- #
 # ------------------------------------ Data loading --------------------------------- #
 # ----------------------------------------------------------------------------------- #
@@ -19,7 +21,7 @@ def __init__(
         batch_size: int,
         max_episodes: int,
         max_steps: int,
-        device: torch.device,
+        device: th.DeviceType,
         split_episodes: bool = False,
         data_mode: str = "all",
         val_ratio: float = 0.1,
@@ -71,7 +73,7 @@ def _load_dataset_from_file(
         filename: str,
         episode_length: int,
         max_episodes: int,
-        device: torch.device,
+        device: th.DeviceType,
         split_episodes: bool = False,
     ) -> Dict[str, torch.Tensor]:
         # Load all episode data
@@ -148,7 +150,7 @@ def _get_tensor_splits(
 
     @staticmethod
     def _load_tactile_sdf_from_file(
-        filename: str, device: torch.device
+        filename: str, device: th.DeviceType
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
 
         with open(filename) as f: