Merge pull request #568 from asi1024/fix-mypy-failure

Fix mypy failures for torch 1.12

Author: xuzijian629

pytorch_pfn_extras/nn/modules/lazy.py

@@ -129,7 +129,7 @@ def _lazy_load_hook(  # type: ignore[no-untyped-def]
 
 class UninitializedParameter(torch.nn.Parameter):
 
-    def __repr__(self) -> str:
+    def __repr__(self) -> str:  # type: ignore[override]
         return 'Uninitialized lazy parameter'
 
     def share_memory_(self) -> 'UninitializedParameter':

pytorch_pfn_extras/nn/parallel/distributed.py

@@ -233,10 +233,10 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
 
     def load_state_dict(
             self,
-            state_dict: 'OrderedDict[str, torch.Tensor]',
+            state_dict: 'Mapping[str, torch.Tensor]',
             strict: bool = True,
     ) -> None:
-        self.module.load_state_dict(state_dict, strict=strict)
+        self.module.load_state_dict(state_dict, strict=strict)  # type: ignore[arg-type]
 
     T_destination = TypeVar('T_destination', bound=Mapping[str, torch.Tensor])
 

pytorch_pfn_extras/onnx/_constants.py

@@ -5,7 +5,7 @@
 if pytorch_pfn_extras.requires("1.12.0"):
     from torch.onnx._constants import *  # NOQA
 else:
-    onnx_default_opset = torch.onnx.symbolic_helper._default_onnx_opset_version
-    onnx_main_opset = torch.onnx.symbolic_helper._onnx_main_opset
-    onnx_stable_opsets = torch.onnx.symbolic_helper._onnx_stable_opsets
+    onnx_default_opset = torch.onnx.symbolic_helper._default_onnx_opset_version  # type: ignore
+    onnx_main_opset = torch.onnx.symbolic_helper._onnx_main_opset  # type: ignore
+    onnx_stable_opsets = torch.onnx.symbolic_helper._onnx_stable_opsets  # type: ignore
     onnx_constant_folding_opsets = torch.onnx.symbolic_helper._constant_folding_opset_versions if pytorch_pfn_extras.requires("1.11.0") else torch.onnx.constant_folding_opset_versions  # type: ignore[attr-defined]

pytorch_pfn_extras/onnx/_globals.py

@@ -1,29 +1,30 @@
 import pytorch_pfn_extras
 import torch
-import torch.onnx.symbolic_helper
 from typing import Optional
 
 
-class _InternalGlobalsBeforeTorch1_11:
-    @property
-    def export_onnx_opset_version(self) -> int:
-        return torch.onnx.symbolic_helper._export_onnx_opset_version
+if pytorch_pfn_extras.requires("1.12.0"):
+    import torch.onnx._globals
+    GLOBALS = torch.onnx._globals.GLOBALS
 
-    @property
-    def operator_export_type(self) -> Optional[torch._C._onnx.OperatorExportTypes]:
-        return torch.onnx.symbolic_helper._operator_export_type
+else:
+    import torch.onnx.symbolic_helper as symhel
 
-    @property
-    def training_mode(self) -> Optional[torch._C._onnx.TrainingMode]:
-        return torch.onnx.symbolic_helper._training_mode
+    class _InternalGlobalsBeforeTorch1_11:
+        @property
+        def export_onnx_opset_version(self) -> int:
+            return symhel._export_onnx_opset_version  # type: ignore
 
-    @property
-    def onnx_shape_inference(self) -> bool:
-        return torch.onnx.symbolic_helper._onnx_shape_inference
+        @property
+        def operator_export_type(self) -> Optional[torch._C._onnx.OperatorExportTypes]:
+            return symhel._operator_export_type  # type: ignore
 
+        @property
+        def training_mode(self) -> Optional[torch._C._onnx.TrainingMode]:
+            return symhel._training_mode  # type: ignore
 
-if pytorch_pfn_extras.requires("1.12.0"):
-    import torch.onnx._globals
-    GLOBALS = torch.onnx._globals.GLOBALS
-else:
-    GLOBALS = _InternalGlobalsBeforeTorch1_11()
+        @property
+        def onnx_shape_inference(self) -> bool:
+            return symhel._onnx_shape_inference  # type: ignore
+
+    GLOBALS = _InternalGlobalsBeforeTorch1_11()  # type: ignore[assignment]

pytorch_pfn_extras/onnx/export_testcase.py

@@ -111,7 +111,7 @@ def _export_util(
         operator_export_type = OperatorExportTypes.ONNX_ATEN if\
             aten else OperatorExportTypes.RAW  # type: ignore
     elif operator_export_type is None:
-        if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE:
+        if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE:  # type: ignore[attr-defined]
             operator_export_type = OperatorExportTypes.ONNX_ATEN_FALLBACK
         else:
             operator_export_type = OperatorExportTypes.ONNX