[Reland] Update mypy to 1.4.1 (pytorch#105227)

This PR re-lands
- [Typing] Fix PEP 484 Violation (pytorch#105022)
- Update mypy to 1.4.1 (pytorch#91983)

That were reverted due to the conflict with internal source repo.

Mostly fixes for PEP-484 violation (i.e. when default arg is set to None, but type is not annotated as optional)
Plus few real fixes:
  - Add missing `_get_upgraders_entry_map` to `torch/_C/__init__.pyi`
  - Add missing return statement to `torch._export. deserialize_graph`
  - Fix error message in `torch.ao.ns.fx.weight_utils.get_lstm_mod_weights`
  - Add assert it `torch/optim/optimizer.py` that Optional list is not None
TODO (in followup PR):
  - Fix erroneous `isinstance` check in `torch/ao/quantization/_pt2e/qat_utils.py`

Unrelated, to bypass CI failures due to the gcc9 dependency update in Ubuntu-18.04:
- Add hack to squash older libstdc++ from conda environment in favor one from OS to `.ci/docker/install_conda.sh`
- Update bazel cuda builds to focal, as with libstdc++-6.0.32 bazel builds loose the ability to catch exceptions (probably because they link with cupti statically, but I could not found where it is done)
Pull Request resolved: pytorch#105227
Approved by: https://github.com/atalman, https://github.com/albanD, https://github.com/Skylion007

.ci/docker/build.sh

@@ -160,6 +160,20 @@ case "$image" in
     TRITON=yes
     INDUCTOR_BENCHMARKS=yes
     ;;
+  pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9)
+    CUDA_VERSION=12.1.0
+    CUDNN_VERSION=8
+    ANACONDA_PYTHON_VERSION=3.10
+    GCC_VERSION=9
+    PROTOBUF=yes
+    DB=yes
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    CONDA_CMAKE=yes
+    TRITON=yes
+    ;;
   pytorch-linux-focal-py3-clang7-asan)
     ANACONDA_PYTHON_VERSION=3.9
     CLANG_VERSION=7

.ci/docker/common/install_conda.sh

@@ -105,5 +105,13 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
     pip_install -r /opt/conda/requirements-docs.txt
   fi
 
+  # HACK HACK HACK
+  # gcc-9 for ubuntu-18.04 from http://ppa.launchpad.net/ubuntu-toolchain-r/test/ubuntu
+  # Pulls llibstdc++6 13.1.0-8ubuntu1~18.04 which is too new for conda
+  # So remove libstdc++6.so.3.29 installed by https://anaconda.org/anaconda/libstdcxx-ng/files?version=11.2.0
+  if grep 18.04.6 /etc/issue >/dev/null; then
+    rm /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/lib/libstdc++.so.6
+  fi
+
   popd
 fi

.ci/docker/requirements-ci.txt

@@ -75,10 +75,10 @@ librosa>=0.6.2 ; python_version < "3.11"
 #Pinned versions:
 #test that import:
 
-mypy==0.960
+mypy==1.4.1
 # Pin MyPy version because new errors are likely to appear with each release
 #Description: linter
-#Pinned versions: 0.960
+#Pinned versions: 1.4.1
 #test that import: test_typing.py, test_type_hints.py
 
 networkx==2.8.8

.github/workflows/docker-builds.yml

@@ -40,6 +40,7 @@ jobs:
           - docker-image-name: pytorch-linux-bionic-cuda11.8-cudnn8-py3-gcc7-inductor-benchmarks
           - docker-image-name: pytorch-linux-bionic-py3.8-clang9
           - docker-image-name: pytorch-linux-bionic-py3.11-clang9
+          - docker-image-name: pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9
           - docker-image-name: pytorch-linux-focal-rocm-n-1-py3
           - docker-image-name: pytorch-linux-focal-rocm-n-py3
           - docker-image-name: pytorch-linux-jammy-cuda11.8-cudnn8-py3.8-clang12

.github/workflows/pull.yml

@@ -305,12 +305,12 @@ jobs:
           { config: "default", shard: 1, num_shards: 1, runner: "linux.4xlarge" },
         ]}
 
-  linux-bionic-cuda12_1-py3_10-gcc9-bazel-test:
-    name: linux-bionic-cuda12.1-py3.10-gcc9-bazel-test
+  linux-focal-cuda12_1-py3_10-gcc9-bazel-test:
+    name: linux-focal-cuda12.1-py3.10-gcc9-bazel-test
     uses: ./.github/workflows/_bazel-build-test.yml
     with:
-      build-environment: linux-bionic-cuda12.1-py3.10-gcc9-bazel-test
-      docker-image-name: pytorch-linux-bionic-cuda12.1-cudnn8-py3-gcc9
+      build-environment: linux-focal-cuda12.1-py3.10-gcc9-bazel-test
+      docker-image-name: pytorch-linux-focal-cuda12.1-cudnn8-py3-gcc9
       cuda-version: "12.1"
       test-matrix: |
         { include: [

.lintrunner.toml

@@ -152,7 +152,7 @@ init_command = [
     '--dry-run={{DRYRUN}}',
     'numpy==1.24.3',
     'expecttest==0.1.3',
-    'mypy==0.960',
+    'mypy==1.4.1',
     'types-requests==2.27.25',
     'types-PyYAML==6.0.7',
     'types-tabulate==0.8.8',

caffe2/contrib/aten/gen_op.py

@@ -49,7 +49,7 @@
     # use faster C loader if available
     from yaml import CSafeLoader as Loader
 except ImportError:
-    from yaml import SafeLoader as Loader  # type: ignore[misc]
+    from yaml import SafeLoader as Loader  # type: ignore[assignment, misc]
 
 
 def write(filename, s):

caffe2/contrib/tensorboard/tensorboard.py

@@ -106,7 +106,7 @@ def graph_def_to_event(step, graph_def):
         wall_time=step, step=step, graph_def=graph_def.SerializeToString())
 
 
-@cli.command("tensorboard-graphs")
+@cli.command("tensorboard-graphs")  # type: ignore[arg-type, attr-defined]
 @click.option("--c2-netdef", type=click.Path(exists=True, dir_okay=False),
               multiple=True)
 @click.option("--tf-dir", type=click.Path(exists=True))
@@ -129,7 +129,7 @@ def parse_net_def(path):
     log.info("Wrote %s graphs to logdir %s", len(events), tf_dir)
 
 
-@cli.command("tensorboard-events")
+@cli.command("tensorboard-events")  # type: ignore[arg-type, attr-defined]
 @click.option("--c2-dir", type=click.Path(exists=True, file_okay=False),
               help="Root directory of the Caffe2 run")
 @click.option("--tf-dir", type=click.Path(writable=True),
@@ -209,4 +209,4 @@ def event(step, values):
 
 
 if __name__ == "__main__":
-    cli()
+    cli()  # type: ignore[misc]

test/test_futures.py

@@ -21,23 +21,23 @@ def test_set_exception(self) -> None:
         error_msg = "Intentional Value Error"
         value_error = ValueError(error_msg)
 
-        f = Future[T]()
+        f = Future[T]()  # type: ignore[valid-type]
         # Set exception
         f.set_exception(value_error)
         # Exception should throw on wait
         with self.assertRaisesRegex(ValueError, "Intentional"):
             f.wait()
 
         # Exception should also throw on value
-        f = Future[T]()
+        f = Future[T]()  # type: ignore[valid-type]
         f.set_exception(value_error)
         with self.assertRaisesRegex(ValueError, "Intentional"):
             f.value()
 
         def cb(fut):
             fut.value()
 
-        f = Future[T]()
+        f = Future[T]()  # type: ignore[valid-type]
         f.set_exception(value_error)
 
         with self.assertRaisesRegex(RuntimeError, "Got the following error"):
@@ -54,7 +54,7 @@ def wait_future(f):
             with self.assertRaisesRegex(ValueError, "Intentional"):
                 f.wait()
 
-        f = Future[T]()
+        f = Future[T]()  # type: ignore[valid-type]
         t = threading.Thread(target=wait_future, args=(f, ))
         t.start()
         f.set_exception(value_error)
@@ -68,7 +68,7 @@ def then_future(f):
             with self.assertRaisesRegex(RuntimeError, "Got the following error"):
                 fut.wait()
 
-        f = Future[T]()
+        f = Future[T]()  # type: ignore[valid-type]
         t = threading.Thread(target=then_future, args=(f, ))
         t.start()
         f.set_exception(value_error)

test/test_torch.py

@@ -29,7 +29,7 @@
 from functools import partial
 from torch import multiprocessing as mp
 from torch.testing import make_tensor
-from torch.testing._internal.common_utils import (
+from torch.testing._internal.common_utils import (  # type: ignore[attr-defined]
     TEST_WITH_TORCHINDUCTOR, TestCase, TEST_WITH_ROCM, run_tests, IS_JETSON,
     IS_WINDOWS, IS_FILESYSTEM_UTF8_ENCODING, NO_MULTIPROCESSING_SPAWN,
     IS_SANDCASTLE, IS_FBCODE, IS_REMOTE_GPU, skipIfTorchInductor, load_tests, slowTest,
@@ -8490,7 +8490,7 @@ def _spawn_method(self, method, arg):
             except RuntimeError:
                 pass
             with mp.Pool(1) as pool:
-                out: list = pool.map(method, [arg])
+                out = pool.map(method, [arg])
                 self.assertTrue(out[0])
 
         def _test_multinomial_invalid_probs(probs):

tools/code_coverage/oss_coverage.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
 import time
 
-from package.oss.cov_json import get_json_report
-from package.oss.init import initialization
-from package.tool.summarize_jsons import summarize_jsons
-from package.util.setting import TestPlatform
-from package.util.utils import print_time
+from package.oss.cov_json import get_json_report  # type: ignore[import]
+from package.oss.init import initialization  # type: ignore[import]
+from package.tool.summarize_jsons import summarize_jsons  # type: ignore[import]
+from package.util.setting import TestPlatform  # type: ignore[import]
+from package.util.utils import print_time  # type: ignore[import]
 
 
 def report_coverage() -> None:

tools/code_coverage/package/tool/summarize_jsons.py

@@ -45,7 +45,7 @@ def transform_file_name(
                 return file_path[file_path.find(folder) :]
     # remove pytorch base folder path
     if platform == TestPlatform.OSS:
-        from package.oss.utils import get_pytorch_folder
+        from package.oss.utils import get_pytorch_folder  # type: ignore[import]
 
         pytorch_foler = get_pytorch_folder()
         assert file_path.startswith(pytorch_foler)

tools/code_coverage/package/util/utils.py

@@ -89,7 +89,9 @@ def get_raw_profiles_folder() -> str:
 
 def detect_compiler_type(platform: TestPlatform) -> CompilerType:
     if platform == TestPlatform.OSS:
-        from package.oss.utils import detect_compiler_type  # type: ignore[misc]
+        from package.oss.utils import (  # type: ignore[assignment, import, misc]
+            detect_compiler_type,
+        )
 
         cov_type = detect_compiler_type()  # type: ignore[call-arg]
     else:
@@ -100,7 +102,7 @@ def detect_compiler_type(platform: TestPlatform) -> CompilerType:
         cov_type = detect_compiler_type()
 
     check_compiler_type(cov_type)
-    return cov_type
+    return cov_type  # type: ignore[no-any-return]
 
 
 def get_test_name_from_whole_path(path: str) -> str:

tools/gen_vulkan_spv.py

@@ -20,7 +20,7 @@
 try:
     from yaml import CLoader as Loader
 except ImportError:
-    from yaml import Loader  # type: ignore[misc]
+    from yaml import Loader  # type: ignore[assignment, misc]
 
 H_NAME = "spv.h"
 CPP_NAME = "spv.cpp"

tools/linter/adapters/workflow_consistency_linter.py

@@ -16,7 +16,7 @@
 try:
     from yaml import CSafeLoader as Loader
 except ImportError:
-    from yaml import SafeLoader as Loader  # type: ignore[misc]
+    from yaml import SafeLoader as Loader  # type: ignore[assignment, misc]
 
 
 class LintSeverity(str, Enum):

tools/lite_interpreter/gen_selected_mobile_ops_header.py

@@ -11,7 +11,7 @@
 try:
     from yaml import CSafeLoader as Loader
 except ImportError:
-    from yaml import SafeLoader as Loader  # type: ignore[misc]
+    from yaml import SafeLoader as Loader  # type: ignore[assignment, misc]
 
 
 if_condition_template_str = """if (kernel_tag_sv.compare("$kernel_tag_name") == 0) {

tools/setup_helpers/generate_code.py

@@ -10,7 +10,7 @@
     # use faster C loader if available
     from yaml import CSafeLoader as YamlLoader
 except ImportError:
-    from yaml import SafeLoader as YamlLoader  # type: ignore[misc]
+    from yaml import SafeLoader as YamlLoader  # type: ignore[assignment, misc]
 
 NATIVE_FUNCTIONS_PATH = "aten/src/ATen/native/native_functions.yaml"
 TAGS_PATH = "aten/src/ATen/native/tags.yaml"

torch/_C/__init__.pyi.in

@@ -429,6 +429,7 @@ def _jit_is_script_object(obj: Any) -> _bool: ...
 def _last_executed_optimized_graph() -> Graph: ...
 def parse_type_comment(comment: str) -> Decl: ...
 def _get_upgraders_map_size() -> _int: ...
+def _get_upgraders_entry_map() -> Dict[str, str]: ...
 def _dump_upgraders_map() -> Dict[str, str]: ...
 def _test_only_populate_upgraders(content: Dict[str, str]) -> None: ...
 def _test_only_remove_upgraders(content: Dict[str, str]) -> None: ...

torch/__init__.py

@@ -417,7 +417,7 @@ def sym_int(a):
     if isinstance(a, SymInt):
         return a
     elif isinstance(a, SymFloat):
-        return math.floor(a) if a >= 0 else math.ceil(a)  # type: ignore[arg-type]
+        return math.floor(a) if a >= 0 else math.ceil(a)  # type: ignore[arg-type, call-overload]
     return py_int(a)  # type: ignore[operator]
 
 def sym_max(a, b):
@@ -1320,7 +1320,7 @@ def manager_path():
     # Some type signatures pulled in from _VariableFunctions here clash with
     # signatures already imported. For now these clashes are ignored; see
     # PR #43339 for details.
-    from torch._C._VariableFunctions import *  # type: ignore[misc] # noqa: F403
+    from torch._C._VariableFunctions import *  # type: ignore[assignment, misc] # noqa: F403
     # Fixup segment_reduce visibility
     _segment_reduce = segment_reduce
     del segment_reduce

torch/_export/serde/schema.py

@@ -10,20 +10,20 @@ class _Union:
     @classmethod
     def create(cls, **kwargs):
         assert len(kwargs) == 1
-        return cls(**{**{f.name: None for f in fields(cls)}, **kwargs})
+        return cls(**{**{f.name: None for f in fields(cls)}, **kwargs})  # type: ignore[arg-type]
 
     def __post_init__(self):
-        assert sum(1 for f in fields(self) if getattr(self, f.name) is not None) == 1
+        assert sum(1 for f in fields(self) if getattr(self, f.name) is not None) == 1  # type: ignore[arg-type, misc]
 
     @property
     def value(self):
-        val = next((getattr(self, f.name) for f in fields(self) if getattr(self, f.name) is not None), None)
+        val = next((getattr(self, f.name) for f in fields(self) if getattr(self, f.name) is not None), None)  # type: ignore[arg-type]
         assert val is not None
         return val
 
     @property
     def type(self):
-        val_type = next((f.name for f in fields(self) if getattr(self, f.name) is not None), None)
+        val_type = next((f.name for f in fields(self) if getattr(self, f.name) is not None), None)  # type: ignore[arg-type]
         assert val_type is not None
         return val_type
 

torch/_export/serde/serialize.py

@@ -9,7 +9,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Union
+from typing import cast, Any, Callable, Dict, Iterator, List, Optional, Tuple, Union
 
 import sympy
 
@@ -749,7 +749,7 @@ def __init__(self):
         self.module = torch.nn.Module()
 
     @contextmanager
-    def save_graph_module(self) -> None:
+    def save_graph_module(self) -> Iterator[None]:
         saved = self.graph, self.module, self.serialized_name_to_node, self.serialized_name_to_meta
         self.graph = torch.fx.Graph()
         self.module = torch.nn.Module()
@@ -773,7 +773,7 @@ def deserialize_sym_int(self, s: SymInt) -> Union[int, torch.SymInt]:
 
                     if vr := self.symbol_name_to_range.get(val.expr_str):
                         symbolic_shapes._constrain_symbol_range(
-                            self.shape_env, sym, vr.lower, vr.upper
+                            self.shape_env, sym, vr.lower, vr.upper  # type: ignore[arg-type]
                         )
 
             return self.shape_env.create_symintnode(sym, hint=val.hint)
@@ -855,6 +855,7 @@ def deserialize_graph(self, serialized_graph: Graph) -> torch.fx.Graph:
         output_node.meta["val"] = tuple(
             arg.meta["val"] for arg in output_node.args[0]
         )
+        return output_node
 
     def deserialize_node(self, serialized_node: Node, target: Callable) -> None:
         if target.__module__ == "_operator":  # TODO(zhxchen17) Follow up on this.
@@ -1050,7 +1051,7 @@ def deserialize_multiple_outputs(self, serialized_node: Node, fx_node: torch.fx.
         self.serialized_name_to_node[fx_node.name] = fx_node
 
     def deserialize_metadata(self, metadata: Dict[str, str]) -> Dict[str, Any]:
-        ret = {}
+        ret: Dict[str, Any] = {}
         if stack_trace := metadata.get("stack_trace"):
             ret["stack_trace"] = stack_trace
 

torch/_export/serde/upgrade.py

@@ -33,7 +33,7 @@ def get_upgraders() -> Dict[str, Tuple[str, str]]:
     """Getting upgraders entry map and operator version map and merge them into one dict."""
     upgraders = torch._C._get_upgraders_entry_map()
     op_version_map = torch._C._get_operator_version_map()
-    output = defaultdict(tuple)
+    output: Dict[str, Tuple[str, str]] = defaultdict(tuple)  # type: ignore[arg-type]
     for opname, entry_list in op_version_map.items():
         if not entry_list:
             raise RuntimeError(f"Op version map has an empty entry for opname {opname}")

torch/_functorch/functional_call.py

@@ -1,5 +1,5 @@
 from collections import Counter
-from typing import Any, Dict, List, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -12,7 +12,7 @@ def functional_call(
     module: "torch.nn.Module",
     parameter_and_buffer_dicts: Union[Dict[str, Tensor], Sequence[Dict[str, Tensor]]],
     args: Union[Any, Tuple],
-    kwargs: Dict[str, Any] = None,
+    kwargs: Optional[Dict[str, Any]] = None,
     *,
     tie_weights: bool = True,
     strict: bool = False,