Activate mypy in ignite.distributed

ignite/distributed/comp_models/base.py

@@ -1,7 +1,7 @@
 import warnings
 from abc import ABCMeta, abstractmethod
 from numbers import Number
-from typing import Any, Callable, List, Optional, Union, cast, overload
+from typing import Any, Callable, List, Optional, Union, cast
 
 import torch
 
@@ -134,20 +134,9 @@ def _apply_op(
             return tensor.to(device=tensor_device)
         return tensor
 
-    @overload
-    def _collective_op(
-        self, tensor: Union[torch.Tensor, Number], fn: Callable, *args: Any, **kwargs: Any
-    ) -> Union[torch.Tensor, Number]:
-        ...
-
-    @overload
     def _collective_op(
         self, tensor: Union[torch.Tensor, Number, str], fn: Callable, *args: Any, **kwargs: Any
     ) -> Union[torch.Tensor, Number, List[str]]:
-        ...
-
-    # mypy doesn't support overload for no-untyped-def check
-    def _collective_op(self, tensor, fn, *args, **kwargs):  # type: ignore
         tensor_to_number = tensor_to_str = False
         device = self.device()
         if isinstance(tensor, Number):
@@ -160,7 +149,7 @@ def _collective_op(self, tensor, fn, *args, **kwargs):  # type: ignore
         tensor = self._apply_op(tensor, device, fn, *args, **kwargs)
 
         if tensor_to_number and tensor.numel() == 1:
-            return tensor.item()
+            return cast(Number, tensor.item())
         elif tensor_to_str:
             return self._decode_str(tensor)
         return tensor
@@ -169,23 +158,15 @@ def all_reduce(self, tensor: Union[torch.Tensor, Number], op: str = "sum") -> Un
         if not isinstance(tensor, (torch.Tensor, Number)):
             raise TypeError("Unhandled input type {}".format(type(tensor)))
 
-        return self._collective_op(tensor, self._do_all_reduce, op)
+        return cast(Union[torch.Tensor, Number], self._collective_op(tensor, self._do_all_reduce, op))
 
     def all_gather(self, tensor: Union[torch.Tensor, Number, str]) -> Union[torch.Tensor, Number, List[str]]:
         if not isinstance(tensor, (torch.Tensor, Number, str)):
             raise TypeError("Unhandled input type {}".format(type(tensor)))
 
         return self._collective_op(tensor, self._do_all_gather)
 
-    @overload
-    def broadcast(self, tensor: Union[torch.Tensor, Number], src: int = 0) -> Union[torch.Tensor, Number]:
-        ...
-
-    @overload
     def broadcast(self, tensor: Union[torch.Tensor, Number, str], src: int = 0) -> Union[torch.Tensor, Number, str]:
-        ...
-
-    def broadcast(self, tensor, src=0):  # type: ignore # mypy doesn't support overload for no-untyped-def check
         if not isinstance(tensor, (torch.Tensor, Number, str)):
             raise TypeError("Unhandled input type {}".format(type(tensor)))
 
@@ -210,7 +191,7 @@ def broadcast(self, tensor, src=0):  # type: ignore # mypy doesn't support overl
         tensor = self._apply_op(tensor, device, self._do_broadcast, src)
 
         if tensor_to_number:
-            return tensor.item()
+            return cast(Number, tensor.item())
         if tensor_to_str:
             list_str = self._decode_str(tensor)
             return list_str[0]
@@ -293,15 +274,7 @@ def all_reduce(self, tensor: Union[torch.Tensor, Number], op: str = "sum") -> Un
     def all_gather(self, tensor: Union[torch.Tensor, Number, str]) -> Union[torch.Tensor, Number, List[str]]:
         return cast(Union[torch.Tensor, Number], tensor)
 
-    @overload
-    def broadcast(self, tensor: Union[torch.Tensor, Number], src: int = 0) -> Union[torch.Tensor, Number]:
-        ...
-
-    @overload
     def broadcast(self, tensor: Union[torch.Tensor, Number, str], src: int = 0) -> Union[torch.Tensor, Number, str]:
-        ...
-
-    def broadcast(self, tensor, src=0):  # type: ignore # mypy doesn't support overload for no-untyped-def check
         return tensor
 
     def _do_all_reduce(self, tensor: torch.Tensor, op: str = "sum") -> torch.Tensor:

ignite/distributed/comp_models/native.py

@@ -2,7 +2,7 @@
 import subprocess
 import warnings
 from distutils.version import LooseVersion
-from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, overload
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple
 
 import torch
 import torch.distributed as dist
@@ -245,7 +245,7 @@ def _dist_worker_task_fn(
             node_rank: int,
             master_addr: str,
             master_port: str,
-            **kw: Any,
+            kw: Any,
         ) -> None:
             from ignite.distributed.utils import _set_model, finalize