Typecheck 42% of exir directory

Author: mergennachin

.lintrunner.toml

@@ -314,57 +314,37 @@ include_patterns = [
     'docs/**/*.py',
     # '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 directory - organized by subdirectories with blacklist exclusions
+    'exir/_serialize/**/*.py',
+    'exir/backend/**/*.py', 
+    'exir/capture/**/*.py',
+    'exir/dialects/**/*.py',
+    'exir/emit/**/*.py',
+    'exir/operator/**/*.py',
+    'exir/passes/**/*.py',
+    'exir/program/**/*.py',
+    'exir/serde/**/*.py',
+    'exir/verification/**/*.py',
+    # exir root level files (sorted alphabetically)
+    'exir/__init__.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',
+    'exir/dim_order_utils.py',
+    'exir/dynamic_shape.py',
+    'exir/error.py',
+    'exir/graph.py',
+    'exir/graph_module.py',
+    'exir/memory.py',
+    'exir/pass_manager.py',
+    'exir/print_program.py',
+    'exir/scalar_type.py',
+    'exir/schema.py',
+    'exir/sym_util.py',
+    'exir/types.py',
+    'exir/version.py',
+    'exir/wrap.py',
     # 'extension/**/*.py',
     'kernels/**/*.py',
     'profiler/**/*.py',
@@ -380,6 +360,50 @@ exclude_patterns = [
     'scripts/check_binary_dependencies.py',
     'profiler/test/test_profiler_e2e.py',
     'backends/arm/test/**',
+    # exir exclusions - files with mypy errors
+    'exir/backend/test/**',
+    'exir/backend/utils.py',
+    'exir/capture/_capture.py',
+    'exir/dialects/backend/test/**',
+    'exir/dialects/edge/arg/model.py',
+    'exir/dialects/edge/op/test/**',
+    'exir/dialects/edge/spec/**',
+    'exir/dialects/edge/test/**',
+    'exir/dialects/test/**',
+    'exir/emit/_emitter.py',
+    'exir/emit/test/**',
+    'exir/lowered_backend_module.py',
+    'exir/memory_planning.py',
+    'exir/operator/convert.py',
+    'exir/operator/test/**',
+    'exir/pass_base.py',
+    'exir/passes/__init__.py',
+    'exir/passes/_quant_patterns_and_replacements.py',
+    'exir/passes/const_prop_pass.py',
+    'exir/passes/constant_prop_pass.py',
+    'exir/passes/dynamic_shape_prop_pass.py',
+    'exir/passes/executorch_prim_ops_registry.py',
+    'exir/passes/memory_planning_pass.py',
+    'exir/passes/prune_empty_tensors_pass.py',
+    'exir/passes/quant_fusion_pass.py',
+    'exir/passes/quantize_io_pass.py',
+    'exir/passes/remove_mixed_type_operators.py',
+    'exir/passes/remove_noop_pass.py',
+    'exir/passes/replace_view_copy_with_view_pass.py',
+    'exir/passes/spec_prop_pass.py',
+    'exir/passes/sym_shape_eval_pass.py',
+    'exir/passes/sym_to_tensor_pass.py',
+    'exir/passes/weights_to_outputs_pass.py',
+    'exir/program/_program.py',
+    'exir/program/test/**',
+    'exir/serde/export_serialize.py',
+    'exir/serde/serialize.py',
+    'exir/tensor.py',
+    'exir/tests/**',
+    'exir/tracer.py',
+    'exir/verification/arg_validator.py',
+    'exir/verification/test/**',
+    'exir/_serialize/test/**',
 ]
 command = [
     'python3',

.mypy.ini

@@ -103,3 +103,6 @@ follow_untyped_imports = True
 
 [mypy-sympy.*]
 ignore_missing_imports = True
+
+[mypy-executorch.exir.verification.bindings]
+ignore_missing_imports = True

exir/_serialize/_dataclass.py

@@ -141,5 +141,5 @@ class Example
         if isinstance(T, enum.EnumMeta):
             data[key] = T[value]
         else:
-            data[key] = T(value)
-    return cls(**data)
+            data[key] = T(value)  # type: ignore[operator]
+    return cls(**data)  # type: ignore[operator]

exir/_serialize/_flatbuffer.py

@@ -193,10 +193,10 @@ def _run_flatc(args: Sequence[str]) -> None:
             subprocess.run([flatc_path] + list(args), check=True)
     else:
         # Expect the `flatc` tool to be on the system path or set as an env var.
-        flatc_path = os.getenv("FLATC_EXECUTABLE")
-        if not flatc_path:
-            flatc_path = "flatc"
-        subprocess.run([flatc_path] + list(args), check=True)
+        flatc_executable = os.getenv("FLATC_EXECUTABLE")
+        if not flatc_executable:
+            flatc_executable = "flatc"
+        subprocess.run([flatc_executable] + list(args), check=True)
 
 
 def _flatc_compile(output_dir: str, schema_path: str, json_path: str) -> None:

exir/_serialize/_named_data_store.py

@@ -121,8 +121,8 @@ def _add_named_data_to_map(
             if self.data_hash_to_buffer_idx.get(hashed, -1) != buffer_idx:
                 raise ValueError(
                     f"Duplicate key {key} with different data. "
-                    f"Existing data: {self.buffers[buffer_idx].buffer}. "
-                    f"New data: {data}."
+                    f"Existing data: {self.buffers[buffer_idx].buffer!r}. "
+                    f"New data: {data!r}."  # type: ignore[str-bytes-safe]
                 )
             self.buffers[buffer_idx].alignment = math.lcm(
                 self.buffers[buffer_idx].alignment, alignment

exir/_serialize/_serialize.py

@@ -6,7 +6,7 @@
 
 # pyre-strict
 
-from typing import Dict, Optional, Set, Tuple
+from typing import Dict, List, Optional, Set, Tuple
 
 from executorch.exir._serialize import _serialize_pte_binary
 
@@ -102,10 +102,10 @@ def serialize_for_executorch(
         )
 
     for tag in all_external_tags:
-        buffers = []
+        buffers: List[bytes] = []
         fqn_to_tensor_entry: Dict[str, TensorEntry] = {}
         # pyre-ignore[16]: Undefined attribute: `Optional` has no attribute `get`.
-        fqn_to_index = emitter_output.external_constant_map.get(tag, {})
+        fqn_to_index = emitter_output.external_constant_map.get(tag, {})  # type: ignore[union-attr]
         # Create a TensorEntry for each external tensor.
         for fqn, index in fqn_to_index.items():
             assert fqn in fqn_to_tensor_layout
@@ -118,13 +118,13 @@ def serialize_for_executorch(
         # Extract external data.
         key_to_data: Dict[str, DataEntry] = {}
         # pyre-ignore[16]: Undefined attribute: `Optional` has no attribute `get`.
-        key_to_buffer_index = named_data.external_data.get(tag, {})
+        key_to_buffer_index = named_data.external_data.get(tag, {})  # type: ignore[union-attr]
         for key, index in key_to_buffer_index.items():
             # pyre-ignore[16]: Undefined attribute: `Optional` has no attribute `buffers`.
             key_to_data[key] = DataEntry(
-                len(buffers), named_data.buffers[index].alignment
+                len(buffers), named_data.buffers[index].alignment  # type: ignore[union-attr]
             )
-            buffers.append(named_data.buffers[index].buffer)
+            buffers.append(named_data.buffers[index].buffer)  # type: ignore[union-attr]
 
         # Serialize into PTD file.
         ptd_files[tag] = data_serializer.serialize(

exir/backend/backend_api.py

@@ -123,7 +123,7 @@ def to_backend(
                 compile_specs=compile_specs,
                 named_data_store_output=preprocess_result.data_store_output,
             )
-            lowered_module.meta = {
+            lowered_module.meta = {  # type: ignore[assignment]
                 "debug_handle_map": preprocess_result.debug_handle_map
             }
             return lowered_module
@@ -311,7 +311,7 @@ def _partition_and_lower_one_graph_module(
             is_submodule,
         )
 
-        lowered_submodule = to_backend(
+        lowered_submodule = to_backend(  # type: ignore[call-arg]
             delegation_spec.backend_id,
             submodule_program,
             delegation_spec.compile_specs,
@@ -449,7 +449,7 @@ def _create_partitions_in_graph_module(
     owning_program: ExportedProgram,
     is_submodule: bool,
 ) -> Dict[str, List[torch.fx.Node]]:
-    backend_id_to_submodule_name = {}
+    backend_id_to_submodule_name: Dict[str, List[str]] = {}
     for tag, delegation_spec in partition_result.partition_tags.items():
         # Create partition with nodes containing this tag. There should only be
         # one contained submodule per tag
@@ -517,10 +517,12 @@ def _create_partitions_in_graph_module(
         # in future edits to the graph. As a result, we just keep track of the node's name
         # and at the end we search for this node in our final graph module
         backend_id_to_submodule_name[delegation_spec.backend_id].append(
-            call_module_node.target
+            call_module_node.target  # type: ignore[arg-type]
         )
 
-    created_submodule_nodes = {key: [] for key in backend_id_to_submodule_name.keys()}
+    created_submodule_nodes: Dict[str, List[torch.fx.Node]] = {
+        key: [] for key in backend_id_to_submodule_name.keys()
+    }
     for backend_id, submodule_name in backend_id_to_submodule_name.items():
         for node in tagged_graph_module.graph.nodes:
             if node.op == "call_module" and node.target in submodule_name:
@@ -615,7 +617,7 @@ def lower_all_submodules_to_backend(
                 compile_specs=compile_spec,
                 named_data_store_output=preprocess_result.data_store_output,
             )
-            lowered_module.meta = {
+            lowered_module.meta = {  # type: ignore[assignment]
                 "debug_handle_map": preprocess_result.debug_handle_map,
             }
             is_submodule = call_submodule_node.meta["is_submodule"]
@@ -698,7 +700,7 @@ def to_backend(
     method_to_partitioner = method_edge_program_partitioners.method_to_partitioner
 
     partitioned_and_lowered_exported_programs = {}
-    backend_id_to_method_submodules_map = {}
+    backend_id_to_method_submodules_map: Dict[str, Dict[str, List[torch.fx.Node]]] = {}
     method_to_tagged_exported_program = {}
 
     for method_name, partitioner_instance in method_to_partitioner.items():

exir/backend/canonical_partitioners/config_partitioner.py

@@ -52,10 +52,9 @@ class PartitionerConfig(ABC):
     the specified backend.
     """
 
-    @classmethod
-    @property
+    @property  # type: ignore[misc]
     @abstractmethod
-    def target_name(cls) -> str:
+    def target_name(self) -> str:
         """
         Target name for this partitioner config. When the Config-Based Partitioner
         encounters a node with a matching target name, it uses this config's methods to
@@ -138,7 +137,7 @@ def filter_fn(node: torch.fx.Node) -> bool:
             """
             if node.op != "call_function":
                 return False
-            target_name = format_target_name(node.target.__name__)  # pyre-ignore
+            target_name = format_target_name(node.target.__name__)  # type: ignore[union-attr]
 
             if target_name in self.target_partitioner_configs:
                 config = self.target_partitioner_configs[target_name]

exir/backend/canonical_partitioners/duplicate_constant_node_pass.py

@@ -28,9 +28,9 @@ def _get_attribute_or_constants(
     if maybe_param is not None:
         constant_or_attribute = maybe_param
     elif maybe_buffer is not None:
-        constant_or_attribute = maybe_buffer
+        constant_or_attribute = maybe_buffer  # type: ignore[assignment]
     elif maybe_lifted_tensor is not None:
-        constant_or_attribute = maybe_lifted_tensor
+        constant_or_attribute = maybe_lifted_tensor  # type: ignore[assignment]
     return constant_or_attribute
 
 

exir/capture/_capture.py

@@ -122,10 +122,10 @@ def _capture_legacy_do_not_use(f, args) -> ExirExportedProgram:
                     outputs=[],
                     # pyre-fixme[6]: For 3rd argument expected `TreeSpec` but got
                     #  `Union[Tensor, Module]`.
-                    in_spec=in_spec,
+                    in_spec=in_spec,  # type: ignore[arg-type]
                     # pyre-fixme[6]: For 4th argument expected `TreeSpec` but got
                     #  `Union[Tensor, Module]`.
-                    out_spec=out_spec,
+                    out_spec=out_spec,  # type: ignore[arg-type]
                 ),
             )
         ],
@@ -207,7 +207,7 @@ def capture(  # noqa: C901
             if isinstance(f, MethodType) and isinstance(f.__self__, torch.nn.Module):
                 with patch_forward(f.__self__, f):
                     ep = export(
-                        cast(torch.nn.Module, f.__self__),
+                        f.__self__,  # type: ignore[redundant-cast]
                         args,
                         dynamic_shapes=dynamic_shapes,
                         strict=True,
@@ -272,7 +272,7 @@ def graph_with_interpreter(*args):
                 graph_with_interpreter,
                 remove="mutations_and_views",
             )
-            assert isinstance(functionalized_callable, Callable)
+            assert callable(functionalized_callable)  # type: ignore[arg-type]
 
             if config.enable_dynamic_shape:
                 fake_tensor_mode = FakeTensorMode(
@@ -357,7 +357,7 @@ def convert_to_fake(x):
                     in_spec=in_spec,
                     # pyre-fixme[6]: For 4th argument expected `TreeSpec` but got
                     #  `Union[None, TreeSpec, Tensor, Module]`.
-                    out_spec=out_spec,
+                    out_spec=out_spec,  # type: ignore[arg-type]
                 ),
             )
         ],