Renamed TheseusDevice as DeviceType.

theseus/__init__.py

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

theseus/constants.py

@@ -57,4 +57,4 @@ def __getitem__(self, dtype):
 
 _SE3_HAT_EPS = EPSDict(float32_eps=5e-6, float64_eps=5e-7)
 
-_TheseusDevice = Optional[Union[str, torch.device]]
+DeviceType = Optional[Union[str, torch.device]]

theseus/core/objective.py

@@ -10,7 +10,7 @@
 
 import torch
 
-from theseus.constants import _TheseusDevice
+from theseus.constants import DeviceType
 from theseus.core.theseus_function import TheseusFunction
 from theseus.geometry.manifold import Manifold
 
@@ -56,7 +56,7 @@ def __init__(self, dtype: Optional[torch.dtype] = None):
 
         self._batch_size: Optional[int] = None
 
-        self.device: _TheseusDevice = 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,7 +7,7 @@
 from typing import Optional, Sequence, Union
 
 import torch
-from theseus.constants import _TheseusDevice
+from theseus.constants import DeviceType
 
 
 class Variable:
@@ -108,7 +108,7 @@ 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: _TheseusDevice = None,
+    device: DeviceType = None,
     dtype: Optional[torch.dtype] = None,
     name: Optional[str] = None,
 ) -> Variable:

theseus/embodied/kinematics/kinematics_model.py

@@ -8,7 +8,7 @@
 
 import torch
 
-from theseus.constants import _TheseusDevice
+from theseus.constants import DeviceType
 from theseus.geometry import SE3, LieGroup, Point2, Vector
 
 RobotModelInput = Union[torch.Tensor, Vector]
@@ -37,7 +37,7 @@ def forward_kinematics(self, robot_pose: RobotModelInput) -> Dict[str, LieGroup]
 
 
 class UrdfRobotModel(KinematicsModel):
-    def __init__(self, urdf_path: str, device: _TheseusDevice = 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,7 +8,7 @@
 
 import torch
 
-from theseus.constants import _TheseusDevice
+from theseus.constants import DeviceType
 from theseus.geometry.manifold import Manifold
 
 
@@ -52,7 +52,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: _TheseusDevice = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "LieGroup":
         pass
@@ -63,7 +63,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: _TheseusDevice = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "LieGroup":
         pass

theseus/geometry/point_types.py

@@ -7,7 +7,7 @@
 
 import torch
 
-from theseus.constants import _TheseusDevice
+from theseus.constants import DeviceType
 from .vector import Vector
 
 
@@ -42,7 +42,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: _TheseusDevice = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Point2":
         if len(size) != 1:
@@ -63,7 +63,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: _TheseusDevice = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Point2":
         if len(size) != 1:
@@ -144,7 +144,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: _TheseusDevice = None,
+        device: DeviceType = None,
         requires_grad: bool = False,
     ) -> "Point3":
         if len(size) != 1:
@@ -165,7 +165,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: _TheseusDevice = 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: theseus.constants._TheseusDevice = 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: theseus.constants._TheseusDevice = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SE2":
         if len(size) != 1:

theseus/geometry/se3.py

@@ -44,7 +44,7 @@ def rand(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: theseus.constants._TheseusDevice = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SE3":
         if len(size) != 1:
@@ -72,7 +72,7 @@ def randn(
         *size: int,
         generator: Optional[torch.Generator] = None,
         dtype: Optional[torch.dtype] = None,
-        device: theseus.constants._TheseusDevice = 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: theseus.constants._TheseusDevice = 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: theseus.constants._TheseusDevice = 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: theseus.constants._TheseusDevice = 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: theseus.constants._TheseusDevice = None,
+        device: theseus.constants.DeviceType = None,
         requires_grad: bool = False,
     ) -> "SO3":
         if len(size) != 1:

theseus/geometry/vector.py

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

theseus/optimizer/linear/utils.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 from typing import Tuple, Union
 import torch
-from theseus.constants import _TheseusDevice
+from theseus.constants import DeviceType
 
 
 # See Nocedal and Wright, Numerical Optimization, pp. 260 and 261
@@ -14,7 +14,7 @@ def convert_to_alpha_beta_damping_tensors(
     damping_eps: float,
     ellipsoidal_damping: bool,
     batch_size: int,
-    device: _TheseusDevice,
+    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__, _TheseusDevice
+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) -> _TheseusDevice:
+    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: th.device = "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: th.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

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

theseus/utils/examples/tactile_pose_estimation/models.py

@@ -47,7 +47,10 @@ def init_tactile_model_from_file(model: nn.Module, filename: pathlib.Path):
 # Set some parameters for the cost weights
 class TactileWeightModel(nn.Module):
     def __init__(
-        self, device: th.device, dim: int = 3, wt_init: Optional[torch.Tensor] = None
+        self,
+        device: th.DeviceType,
+        dim: int = 3,
+        wt_init: Optional[torch.Tensor] = None,
     ):
         super().__init__()
 
@@ -63,7 +66,7 @@ def forward(self):
 
 def create_tactile_models(
     model_type: str,
-    device: th.device,
+    device: th.DeviceType,
     measurements_model_path: Optional[pathlib.Path] = None,
 ) -> Tuple[
     Optional[TactileMeasModel],
@@ -119,7 +122,7 @@ def create_tactile_models(
 
 def get_tactile_nn_measurements_inputs(
     batch: Dict[str, torch.Tensor],
-    device: th.device,
+    device: th.DeviceType,
     class_label: int,
     num_classes: int,
     min_win_mf: int,
@@ -205,7 +208,7 @@ def get_tactile_nn_measurements_inputs(
 
 
 def get_tactile_motion_capture_inputs(
-    batch: Dict[str, torch.Tensor], device: th.device, time_steps: int
+    batch: Dict[str, torch.Tensor], device: th.DeviceType, time_steps: int
 ):
     inputs = {}
     captures = batch["eff_poses"].to(device)
@@ -225,7 +228,7 @@ def get_tactile_cost_weight_inputs(qsp_model, mf_between_model):
 
 def get_tactile_initial_optim_vars(
     batch: Dict[str, torch.Tensor],
-    device: th.device,
+    device: th.DeviceType,
     time_steps: int,
 ):
     inputs = {}
@@ -244,7 +247,7 @@ def update_tactile_pushing_inputs(
     measurements_model: nn.Module,
     qsp_model: nn.Module,
     mf_between_model: nn.Module,
-    device: th.device,
+    device: th.DeviceType,
     cfg: omegaconf.DictConfig,
     theseus_inputs: Dict[str, torch.Tensor],
 ):

theseus/utils/examples/tactile_pose_estimation/pose_estimator.py

@@ -17,7 +17,7 @@ def __init__(
         max_window_moving_frame: int,
         step_window_moving_frame: int,
         rectangle_shape: Tuple[float, float],
-        device: th.device,
+        device: th.DeviceType,
         optimizer_cls: Optional[Type[th.NonlinearLeastSquares]] = LevenbergMarquardt,
         max_iterations: int = 3,
         step_size: float = 1.0,