improve annotation device parameters where a device ordinal is allowed (#113647)

nairbv · pytorchmergebot · commit dbb96ef30da4 · 2023-11-17T14:41:22.000Z
Using mypy in code that depends on pytorch, I noticed that the type annotation doesn't allow a device ordinal. `error: Argument "device" to "to_empty" of "Module" has incompatible type "int"; expected "str | device" [arg-type]` Pull Request resolved: #113647 Approved by: https://github.com/albanD
diff --git a/tools/pyi/gen_pyi.py b/tools/pyi/gen_pyi.py
@@ -74,7 +74,7 @@ def should_bind_method(python_func: PythonSignatureNativeFunctionPair) -> bool:
 # TODO: Consider defining some aliases for our Union[...] types, to make
 # the stubs to read on the human eye.
 
-DEVICE_PARAM = "device: Device = None"
+DEVICE_PARAM = "device: Optional[DeviceLikeType] = None"
 FACTORY_PARAMS = f"dtype: Optional[_dtype] = None, {DEVICE_PARAM}, requires_grad: _bool = False, pin_memory: _bool = False"
 
 # NOTE: specifying indices for Tensor.__getitem__
@@ -624,7 +624,7 @@ def gen_pyi(
                                 "size: Optional[_size] = None",
                                 "*",
                                 "dtype: Optional[_dtype] = None",
-                                "device: Union[_device, str, None] = None",
+                                "device: Optional[DeviceLikeType] = None",
                                 "requires_grad: _bool = False",
                                 "check_invariants: Optional[_bool] = None",
                             ]
@@ -645,7 +645,7 @@ def gen_pyi(
                             "obj: Any",
                             "*",
                             "dtype: Optional[_dtype] = None",
-                            "device: Union[_device, str, None] = None",
+                            "device: Optional[DeviceLikeType] = None",
                             "copy: Optional[_bool] = None",
                             "requires_grad: _bool = False",
                         ]
@@ -662,7 +662,7 @@ def gen_pyi(
                             "dtype: _dtype",
                             "count: int = -1",
                             "offset: int = 0",
-                            "device: Union[_device, str, None] = None",
+                            "device: Optional[DeviceLikeType] = None",
                             "requires_grad: _bool = False",
                         ]
                     )
@@ -700,7 +700,7 @@ def gen_pyi(
                             "size: Optional[_size] = None",
                             "*",
                             "dtype: Optional[_dtype] = None",
-                            "device: Union[_device, str, None] = None",
+                            "device: Optional[DeviceLikeType] = None",
                             "requires_grad: _bool = False",
                             "check_invariants: Optional[_bool] = None",
                             "is_coalesced: Optional[_bool] = None",
@@ -719,7 +719,7 @@ def gen_pyi(
                             "*",
                             "dtype: Optional[_dtype] = None",
                             "layout: Optional[_layout] = None",
-                            "device: Union[_device, str, None] = None",
+                            "device: Optional[DeviceLikeType] = None",
                             "requires_grad: _bool = False",
                             "check_invariants: Optional[_bool] = None",
                         ]
@@ -1100,7 +1100,7 @@ def replace_special_case(hint: str) -> str:
                 )
                 for args in [
                     "dtype: _dtype",
-                    "device: Optional[Union[_device, str]] = None, dtype: Optional[_dtype] = None",
+                    "device: Optional[DeviceLikeType] = None, dtype: Optional[_dtype] = None",
                     "other: Tensor",
                 ]
             ],
diff --git a/torch/_C/_VariableFunctions.pyi.in b/torch/_C/_VariableFunctions.pyi.in
@@ -33,6 +33,8 @@ from torch.types import (
     Number,
 )
 
+from torch._prims_common import DeviceLikeType
+
 ${function_hints}
 
 ${all_directive}
diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in
@@ -51,6 +51,8 @@ from torch.types import (
     Storage,
 )
 
+from torch._prims_common import DeviceLikeType
+
 # This module is defined in torch/csrc/Module.cpp
 
 from . import _functorch, _lazy, _lazy_ts_backend, _nn, _onnx, _VariableFunctions, _cpu
@@ -90,7 +92,7 @@ class device:
 
     # THPDevice_pynew
     @overload
-    def __init__(self, device: Union[_device, _int, str]) -> None: ...
+    def __init__(self, device: DeviceLikeType) -> None: ...
     @overload
     def __init__(self, type: str, index: _int) -> None: ...
 
@@ -398,11 +400,11 @@ def _jit_set_emit_hooks(
 def _jit_get_emit_hooks() -> Tuple[Callable, Callable]: ...
 def _load_for_lite_interpreter(
     filename: Union[str, Path],
-    map_location: Union[_device, str, None],
+    map_location: Optional[DeviceLikeType],
 ): ...
 def _load_for_lite_interpreter_from_buffer(
     buffer: BinaryIO,
-    map_location: Union[_device, str, None],
+    map_location: Optional[DeviceLikeType],
 ): ...
 def _export_operator_list(module: LiteScriptModule): ...
 def _quantize_ondevice_ptq_dynamic(module: LiteScriptModule, method_name: str): ...
@@ -647,20 +649,20 @@ def _parse_source_def(src: str) -> Def: ...
 def import_ir_module(
     cu: CompilationUnit,
     filename: Union[str, Path],
-    map_location: Union[_device, str, None],
+    map_location: Optional[DeviceLikeType],
     extra_files: Dict[str, Any],
 ) -> ScriptModule: ...
 def import_ir_module_from_buffer(
     cu: CompilationUnit,
     buffer: BinaryIO,
-    map_location: Union[_device, str, None],
+    map_location: Optional[DeviceLikeType],
     extra_files: Dict[str, Any],
 ) -> ScriptModule: ...
 def _import_ir_module_from_package(
     cu: CompilationUnit,
     reader: PyTorchFileReader,
     storage_context: DeserializationStorageContext,
-    map_location: Union[_device, str, None],
+    map_location: Optional[DeviceLikeType],
     ts_id: str,
 ) -> ScriptModule: ...
 def _assign_output_shapes(graph: Graph, inputs: List[Tensor]) -> Graph: ...
@@ -1370,7 +1372,7 @@ def _get_privateuse1_backend_name() -> str: ...
 # Defined in torch/csrc/Generator.cpp
 class Generator:
     device: _device
-    def __init__(self, device: Union[_device, str, None] = None) -> None: ...
+    def __init__(self, device: Optional[DeviceLikeType] = None) -> None: ...
     def get_state(self) -> Tensor: ...
     def set_state(self, _new_state: Tensor) -> Generator: ...
     def set_offset(self, offset: _int) -> Generator: ...
diff --git a/torch/_prims_common/__init__.py b/torch/_prims_common/__init__.py
@@ -52,7 +52,7 @@
 FloatLike = (float, torch.SymFloat)
 IntWithoutSymInt = int
 FloatWithoutSymFloat = float
-DeviceLikeType = Union[str, torch.device]
+DeviceLikeType = Union[str, torch.device, int]
 Tensor = torch.Tensor
 
 
@@ -1857,7 +1857,7 @@ def dtype_or_default(dtype: Optional[torch.dtype]) -> torch.dtype:
     return dtype if dtype is not None else torch.get_default_dtype()
 
 
-def device_or_default(device: Optional[torch.device]) -> torch.device:
+def device_or_default(device: Optional[DeviceLikeType]) -> DeviceLikeType:
     return device if device is not None else torch.device("cpu")
 
 
diff --git a/torch/_refs/__init__.py b/torch/_refs/__init__.py
@@ -1939,7 +1939,7 @@ def item(a: TensorLikeType) -> NumberType:
 # fast path when `to` returns an alias to input. This mimics the same function in aten
 def _to_will_alias(
     a: TensorLikeType,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     dtype: Optional[torch.dtype] = None,
     copy: Optional[bool] = None,
     layout: Optional[torch.layout] = None,
@@ -4478,7 +4478,7 @@ def empty(
     *shape,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     requires_grad: bool = False,
     pin_memory: bool = False,
     memory_format: torch.memory_format = torch.contiguous_format,
@@ -4518,7 +4518,7 @@ def empty_permuted(
     physical_layout,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     requires_grad: bool = False,
     pin_memory: bool = False,
 ) -> TensorLikeType:
@@ -4539,7 +4539,7 @@ def new_empty(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     dtype = a.dtype if dtype is None else dtype
@@ -4564,7 +4564,7 @@ def new_empty_strided(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     """
@@ -4591,7 +4591,7 @@ def zeros(
     *size,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -4619,7 +4619,7 @@ def new_zeros(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -4644,7 +4644,7 @@ def ones(
     *size,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -4672,7 +4672,7 @@ def new_ones(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -4700,7 +4700,7 @@ def new_full(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     dtype = a.dtype if dtype is None else dtype
@@ -4723,7 +4723,7 @@ def empty_like(
     a: TensorLikeType,
     *,
     dtype: Optional[torch.dtype] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     layout: Optional[torch.layout] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
@@ -4769,7 +4769,7 @@ def arange(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -4896,7 +4896,7 @@ def linspace(
     steps: NumberType,
     *,
     dtype: Optional[torch.dtype] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     layout: torch.layout = torch.strided,
     pin_memory: bool = False,
     requires_grad: bool = False,
@@ -4986,7 +4986,7 @@ def logspace(
     base: NumberType = 10,
     *,
     dtype: Optional[torch.dtype] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     layout: torch.layout = torch.strided,
     pin_memory: bool = False,
     requires_grad: bool = False,
@@ -5185,7 +5185,7 @@ def empty_strided(
     strides: StrideType,
     *,
     dtype: Optional[torch.dtype] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     layout: torch.layout = torch.strided,
     requires_grad: bool = False,
     pin_memory: bool = False,
@@ -5215,7 +5215,7 @@ def eye(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,  # TODO: unused
 ) -> TensorLikeType:
@@ -5257,7 +5257,7 @@ def full(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
 ) -> TensorLikeType:
@@ -5284,7 +5284,7 @@ def full_like(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -5308,7 +5308,7 @@ def zeros_like(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -5332,7 +5332,7 @@ def ones_like(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: Optional[torch.layout] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
     requires_grad: bool = False,
     memory_format: torch.memory_format = torch.preserve_format,
@@ -5354,7 +5354,7 @@ def ones_like(
 def randn(
     *shape,
     dtype: Optional[torch.dtype] = None,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     layout: Optional[torch.layout] = None,
     requires_grad: bool = False,
     pin_memory: bool = False,
@@ -5381,7 +5381,7 @@ def scalar_tensor(
     *,
     dtype: Optional[torch.dtype] = None,
     layout: torch.layout = torch.strided,
-    device: Optional[torch.device] = None,
+    device: Optional[DeviceLikeType] = None,
     pin_memory: bool = False,
 ) -> TensorLikeType:
     utils.check_layout(layout)
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
@@ -5,6 +5,7 @@
 import weakref
 
 import torch
+from torch._prims_common import DeviceLikeType
 from ..parameter import Parameter
 import torch.utils.hooks as hooks
 
@@ -1016,7 +1017,7 @@ def bfloat16(self: T) -> T:
         """
         return self._apply(lambda t: t.bfloat16() if t.is_floating_point() else t)
 
-    def to_empty(self: T, *, device: Union[str, device], recurse: bool = True) -> T:
+    def to_empty(self: T, *, device: Optional[DeviceLikeType], recurse: bool = True) -> T:
         r"""Move the parameters and buffers to the specified device without copying storage.
 
         Args:
@@ -1031,7 +1032,7 @@ def to_empty(self: T, *, device: Union[str, device], recurse: bool = True) -> T:
         return self._apply(lambda t: torch.empty_like(t, device=device), recurse=recurse)
 
     @overload
-    def to(self, device: Optional[Union[int, str, device]] = ..., dtype: Optional[Union[dtype, str]] = ...,
+    def to(self, device: Optional[DeviceLikeType] = ..., dtype: Optional[Union[dtype, str]] = ...,
            non_blocking: bool = ...) -> Self:
         ...
 
diff --git a/torch/nn/utils/rnn.pyi b/torch/nn/utils/rnn.pyi
@@ -13,7 +13,9 @@ from typing import (
 from typing_extensions import Self
 
 from torch import Tensor
-from torch.types import _device, _dtype
+
+from torch._prims_common import DeviceLikeType
+from torch.types import _dtype
 
 class PackedSequence_(NamedTuple):
     data: Tensor
@@ -54,7 +56,7 @@ class PackedSequence(PackedSequence_):
     @overload
     def to(
         self: _T,
-        device: Optional[Union[_device, str]] = None,
+        device: Optional[DeviceLikeType] = None,
         dtype: Optional[_dtype] = None,
         non_blocking: bool = False,
         copy: bool = False,
diff --git a/torch/types.py b/torch/types.py
diff --git a/torchgen/api/python.py b/torchgen/api/python.py
