Typecheck 25% of exir directory

Author: mergennachin

.lintrunner.toml

@@ -315,6 +315,56 @@ include_patterns = [
     # 'examples/**/*.py',
     'examples/openvino/**/*.py',
     # 'exir/**/*.py',
+    # Phase 1: Start with simplest exir files (Batch 1)
+    'exir/version.py',
+    'exir/scalar_type.py', 
+    'exir/error.py',
+    'exir/_warnings.py',
+    'exir/types.py',
+    # Phase 1: Batch 2 - More utility files
+    'exir/dynamic_shape.py',
+    'exir/memory.py',
+    'exir/dim_order_utils.py', 
+    'exir/wrap.py',
+    # Phase 1: Batch 3 - dialects subdirectory (5 files)
+    'exir/dialects/__init__.py',
+    'exir/dialects/_ops.py',
+    'exir/dialects/backend/_ops.py',
+    'exir/dialects/edge/dtype/supported.py',
+    'exir/dialects/edge/dtype/utils.py',
+    # Phase 1: Batch 3+ - operator utility
+    'exir/operator/util.py',
+    # Phase 1: Batch 4 - More subdirectories (6 files)
+    'exir/program/__init__.py',
+    'exir/program/_fake_program.py',
+    'exir/emit/__init__.py', 
+    'exir/capture/__init__.py',
+    'exir/capture/_config.py',
+    'exir/verification/dev_html.py',
+    # Phase 1: Batch 5 - Fixed problematic files (3 files)
+    'exir/operator/manip.py',
+    'exir/dialects/edge/dtype/runner.py',
+    'exir/serde/schema_check.py',
+    # Phase 1: Batch 6 - Final root-level fixes (3 files)
+    'exir/common.py',
+    'exir/sym_util.py',
+    'exir/graph_module.py',
+    # Phase 1: Batch 7 - Clean files + fixed files (7 files)
+    'exir/schema.py',
+    'exir/print_program.py',
+    'exir/pass_manager.py',
+    'exir/graph.py',
+    'exir/control_flow.py',
+    'exir/delegate.py',
+    'exir/backend/partitioner.py',
+    # Phase 1: Batch 8 - Clean files to reach 25% coverage (7 files)
+    'exir/__init__.py',
+    'exir/capture/_unlift.py', 
+    'exir/serde/__init__.py',
+    'exir/serde/union.py',
+    'exir/serde/schema.py',
+    'exir/_serialize/__init__.py',
+    'exir/_serialize/padding.py',
     # 'extension/**/*.py',
     'kernels/**/*.py',
     'profiler/**/*.py',

.mypy.ini

@@ -100,3 +100,6 @@ ignore_missing_imports = True
 
 [mypy-torchao.*]
 follow_untyped_imports = True
+
+[mypy-sympy.*]
+ignore_missing_imports = True

exir/backend/backend_details.py

@@ -7,16 +7,16 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 from executorch.exir._serialize._named_data_store import NamedDataStoreOutput
 
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from torch.export.exported_program import ExportedProgram
 
 
-def enforcedmethod(func):
-    func.__enforcedmethod__ = True
+def enforcedmethod(func: Callable[..., Any]) -> Callable[..., Any]:
+    func.__enforcedmethod__ = True  # type: ignore[attr-defined]
     return func
 
 

exir/backend/partitioner.py

@@ -59,7 +59,7 @@ class Partitioner(ABC):
     def __init__(
         self,
         spec: Mapping[Union[str, int, float, bool], object] = MappingProxyType({}),
-    ):
+    ) -> None:
         self._spec = spec
 
     def __call__(self, exported_program: ExportedProgram) -> PartitionResult:
@@ -69,7 +69,7 @@ def __call__(self, exported_program: ExportedProgram) -> PartitionResult:
     def spec(self) -> Mapping[Union[str, int, float, bool], object]:
         return self._spec
 
-    @enforcedmethod
+    @enforcedmethod  # type: ignore[misc]
     @abstractmethod
     def partition(self, exported_program: ExportedProgram) -> PartitionResult:
         """

exir/common.py

@@ -104,9 +104,9 @@ def override_logger(
     try:
         oldLevel = logging.root.level
         logging.root.setLevel(newLevel)
+        oldFormatters = []
         if fmtstr:
             newformatter = logging.Formatter(fmtstr, None, "%")
-            oldFormatters = []
             for handler in logging.root.handlers:
                 oldFormatters.append(handler.formatter)
                 handler.formatter = newformatter

exir/control_flow.py

@@ -103,7 +103,7 @@ def _make_submodule(
             f"Expect function '{fn.__name__}' to be decorated with tracing_context.",
         )
     # pyre-ignore
-    args = fn.__tracing_inputs__
+    args = fn.__tracing_inputs__  # type: ignore[attr-defined]
     # TODO(yidi): we don't want to enable here because we are not gonna use this code path in the future anyways
     gm, _ = flattened_dispatch_trace(fn, args, set(), enable_functionalization=False)
     output = next(iter(reversed(gm.graph.nodes)))
@@ -122,7 +122,7 @@ def _make_submodule(
         output.args = tuple(output.args[0])
         gm.recompile()
     # pyre-fixme[16]: `GraphModule` has no attribute `__tracing_inputs__`.
-    gm.__tracing_inputs__ = args
+    gm.__tracing_inputs__ = args  # type: ignore[attr-defined]
     return gm
 
 
@@ -198,7 +198,7 @@ def wrapper(
 
             return f(*args)
 
-        wrapper.__tracing_inputs__ = inputs  # pyre-ignore
+        wrapper.__tracing_inputs__ = inputs  # type: ignore[attr-defined]
         return wrapper
 
     return decorator

exir/delegate.py

@@ -42,9 +42,11 @@
     LOWERED_BACKEND_MODULE_TYPE = "LoweredBackendModule"
 
     # pyre-ignore
-    def trace_call_delegate(proxy_mode, func_overload, lowered_module, *args):
+    def trace_call_delegate(
+        proxy_mode: Any, func_overload: Any, lowered_module: Any, *args: Any
+    ) -> Any:
         # pyre-ignore
-        def _unwrap_proxy(e):
+        def _unwrap_proxy(e: Any) -> Any:
             if not isinstance(e, (torch.Tensor, torch.SymInt, torch.SymFloat)):
                 return e
             return get_proxy_slot(
@@ -151,7 +153,7 @@ def is_lowered_module(obj: Any) -> bool:
     def get_lowered_module_name(
         root: torch.nn.Module,
         # pyre-ignore: Undefined or invalid type [11]: Annotation `LoweredBackendModule` is not defined as a type.
-        lowered_module: LOWERED_BACKEND_MODULE_TYPE,  # noqa
+        lowered_module: Any,  # noqa
     ) -> str:
         """
         Adds the given lowered_module into the given root module and returns the

exir/delegate.pyi

@@ -17,5 +17,5 @@ def is_lowered_module(obj: Any) -> bool: ...
 def get_lowered_module_name(
     root: torch.nn.Module,
     # pyre-ignore: Undefined or invalid type [11]: Annotation `LoweredBackendModule` is not defined as a type.
-    lowered_module: LOWERED_BACKEND_MODULE_TYPE,  # noqa
+    lowered_module: Any,  # noqa
 ) -> str: ...

exir/dialects/edge/dtype/runner.py

@@ -30,7 +30,7 @@ def _get_types(inputs: Dict[str, List[BaseArg]]) -> List[ArgType]:
     @staticmethod
     def _get_args_kwargs(
         inputs: Dict[str, List[BaseArg]],
-        dtypes: Tuple[Optional[torch.dtype]],
+        dtypes: Tuple[Optional[torch.dtype], ...],
         mode: ArgMode,
     ) -> Tuple[List[BaseArg], Dict[str, BaseKwarg]]:
         """Construct args and kwargs for op given dtypes."""
@@ -71,16 +71,20 @@ def run_dtypes(
         self,
         name: str,
         inputs: Dict[str, List[BaseArg]],
-        dtypes: Tuple[Optional[torch.dtype]],
+        dtypes: Tuple[Optional[torch.dtype], ...],
         argmode: ArgMode = ArgMode.RANDOM,
     ) -> Tuple[
-        bool, str, Tuple[Optional[torch.dtype]], List[BaseArg], Dict[str, BaseKwarg]
+        bool,
+        str,
+        Tuple[Optional[torch.dtype], ...],
+        List[BaseArg],
+        Dict[str, BaseKwarg],
     ]:
         args, kwargs = DtypeRunner._get_args_kwargs(inputs, dtypes, argmode)
         op = get_callable(name)
         try:
             res = op(*args, **kwargs)
-            ret_dtypes = ()
+            ret_dtypes: Tuple[torch.dtype, ...] = ()
             if "returns" in inputs:
                 ret_dtypes = tuple(extract_return_dtype(res, inputs["returns"]))
             return (True, name, dtypes + ret_dtypes, args, kwargs)
@@ -112,7 +116,11 @@ def run(
         argmode: ArgMode = ArgMode.ONES,
     ) -> List[
         Tuple[
-            bool, str, Tuple[Optional[torch.dtype]], List[BaseArg], Dict[str, BaseKwarg]
+            bool,
+            str,
+            Tuple[Optional[torch.dtype], ...],
+            List[BaseArg],
+            Dict[str, BaseKwarg],
         ]
     ]:
         results = []

exir/graph_module.py

@@ -42,7 +42,7 @@ def _get_submodule(
     assert submod_node.op == "get_attr"
     assert isinstance(submod_node.target, str)
     submodule = graph_module.get_submodule(submod_node.target)
-    # pyre-ignore
+    assert isinstance(submodule, torch.nn.Module)
     return submod_node.target, submodule, node
 
 
@@ -67,7 +67,7 @@ def get_control_flow_submodules(
         if node.target is torch.ops.higher_order.map_impl:
             control_flow_submodules.append(_get_submodule(graph_module, node, 0))
 
-    return control_flow_submodules
+    return control_flow_submodules  # type: ignore[return-value]
 
 
 def bfs_trace_with_node_process(