Fixed MyPy TypeChecker Issues

.lintrunner.toml

@@ -299,12 +299,14 @@ include_patterns = [
     # TODO(https://github.com/pytorch/executorch/issues/7441): Gradually start enabling all folders.
     # 'backends/**/*.py',
     'backends/arm/**/*.py',
+    'backends/openvino/**/*.py',
     'build/**/*.py',
     'codegen/**/*.py',
     # 'devtools/**/*.py',
     'devtools/visualization/**/*.py',
     'docs/**/*.py',
     # 'examples/**/*.py',
+    'examples/openvino/**/*.py',
     # 'exir/**/*.py',
     # 'extension/**/*.py',
     'kernels/**/*.py',

backends/openvino/__init__.py

@@ -2,4 +2,4 @@
 from .preprocess import OpenvinoBackend
 from .quantizer.quantizer import OpenVINOQuantizer
 
-__all__ = [OpenvinoBackend, OpenvinoPartitioner, OpenVINOQuantizer]
+__all__ = ["OpenvinoBackend", "OpenvinoPartitioner", "OpenVINOQuantizer"]

backends/openvino/partitioner.py

@@ -15,7 +15,9 @@
     PartitionResult,
 )
 from executorch.exir.backend.utils import tag_constant_data
-from openvino.frontend.pytorch.torchdynamo.op_support import OperatorSupport
+from openvino.frontend.pytorch.torchdynamo.op_support import (  # type: ignore[import-untyped]
+    OperatorSupport,
+)
 
 from torch.export.exported_program import ExportedProgram
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
@@ -53,8 +55,11 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
         if node.op != "call_function":
             return False
 
-        options = []
-        op_type = node.target.__name__
+        options: list[str] = []
+        if not isinstance(node.target, str):
+            op_type = node.target.__name__
+        else:
+            op_type = str(node.target)
         supported_ops = OperatorSupport(options)._support_dict
         if op_type == "getitem":
             return True

backends/openvino/preprocess.py

@@ -12,7 +12,9 @@
     PreprocessResult,
 )
 from executorch.exir.backend.compile_spec_schema import CompileSpec
-from openvino.frontend.pytorch.torchdynamo.compile import openvino_compile
+from openvino.frontend.pytorch.torchdynamo.compile import (  # type: ignore[import-untyped]
+    openvino_compile,
+)
 
 
 @final

backends/openvino/quantizer/__init__.py

@@ -1,3 +1,3 @@
 from .quantizer import OpenVINOQuantizer, quantize_model
 
-__all__ = [OpenVINOQuantizer, quantize_model]
+__all__ = ["OpenVINOQuantizer", "quantize_model"]

backends/openvino/quantizer/quantizer.py

@@ -6,16 +6,20 @@
 
 from collections import defaultdict
 from enum import Enum
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, DefaultDict, Dict, List, Optional, Tuple, Type
 
-import nncf
-import nncf.common.quantization as quantization
-import nncf.experimental.torch.fx as nncf_fx
+import nncf  # type: ignore[import-untyped]
+import nncf.common.quantization as quantization  # type: ignore[import-untyped]
+import nncf.experimental.torch.fx as nncf_fx  # type: ignore[import-untyped]
 
 import torch.fx
 
-from nncf.common.graph.graph import NNCFGraph
-from torch.ao.quantization.observer import HistogramObserver, PerChannelMinMaxObserver
+from nncf.common.graph.graph import NNCFGraph  # type: ignore[import-untyped]
+from torch.ao.quantization.observer import (
+    HistogramObserver,
+    PerChannelMinMaxObserver,
+    UniformQuantizationObserverBase,
+)
 from torch.ao.quantization.quantizer.quantizer import (
     EdgeOrNode,
     QuantizationAnnotation,
@@ -117,13 +121,15 @@ def annotate(self, model: torch.fx.GraphModule) -> torch.fx.GraphModule:
         quantization_setup = self.get_nncf_quantization_setup(model, nncf_graph)
 
         graph = model.graph
-        node_vs_torch_annotation = defaultdict(QuantizationAnnotation)
+        node_vs_torch_annotation: DefaultDict[torch.fx.Node, QuantizationAnnotation] = (
+            defaultdict(QuantizationAnnotation)
+        )
 
         for qp in quantization_setup.quantization_points.values():
             edge_or_node, annotation = self._get_edge_or_node_and_annotation(
                 graph, nncf_graph, qp, node_vs_torch_annotation
             )
-            qspec = self._get_torch_ao_qspec_from_qp(qp)
+            qspec: QuantizationSpecBase = self._get_torch_ao_qspec_from_qp(qp)
             self._fill_torch_ao_annotation(edge_or_node, qspec, annotation)
 
         for quantizer_ids in quantization_setup.unified_scale_groups.values():
@@ -199,6 +205,9 @@ def _get_unified_scales_root_quantizer_id(
             ):
                 root_quantizer_id = quantizer_id
                 nncf_node_quantizer_id = nncf_node.node_id
+        if root_quantizer_id is None:
+            msg = "Root quantizer ids can't be None"
+            raise nncf.InternalError(msg)
         return root_quantizer_id
 
     @staticmethod
@@ -299,6 +308,8 @@ def _get_torch_ao_qspec_from_qp(
         qconfig = qp.qconfig
         is_weight = qp.is_weight_quantization_point()
 
+        observer: Type[UniformQuantizationObserverBase]
+
         if qconfig.per_channel:
             torch_qscheme = (
                 torch.per_channel_symmetric

backends/openvino/tests/models/test_classification.py

@@ -1,10 +1,10 @@
-import timm
+import timm  # type: ignore[import-untyped]
 import torch
-import torchvision.models as torchvision_models
+import torchvision.models as torchvision_models  # type: ignore[import-untyped]
 from executorch.backends.openvino.tests.ops.base_openvino_op_test import (
     BaseOpenvinoOpTest,
 )
-from transformers import AutoModel
+from transformers import AutoModel  # type: ignore[import-untyped]
 
 classifier_params = [
     {"model": ["torchvision", "resnet50", (1, 3, 224, 224)]},

backends/openvino/tests/ops/base_openvino_op_test.py

@@ -68,6 +68,7 @@ def execute_layer_test(
             runtime = Runtime.get()
             program = runtime.load_program(exec_prog.buffer)
             method = program.load_method("forward")
+            assert method is not None
             outputs = method.execute(sample_inputs)
 
             # Compare the outputs with the reference outputs

backends/openvino/tests/test_runner.py

@@ -1,12 +1,12 @@
 import argparse
 import unittest
 
-import nncf.torch
+import nncf.torch  # type: ignore[import-untyped]
 
 
 class OpenvinoTestSuite(unittest.TestSuite):
 
-    test_params = {}
+    test_params: dict[str, str] = {}
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -51,7 +51,7 @@ def parse_arguments():
     )
 
     args, ns_args = parser.parse_known_args(namespace=unittest)
-    test_params = {}
+    test_params: dict[str, str] = {}
     test_params["device"] = args.device
     test_params["pattern"] = args.pattern
     test_params["test_type"] = args.test_type

examples/openvino/aot_optimize_and_infer.py

@@ -4,8 +4,11 @@
 # except in compliance with the License. See the license file found in the
 # LICENSE file in the root directory of this source tree.
 
+# mypy: disable-error-code=import-untyped
+
 import argparse
 import time
+from typing import cast, List, Optional
 
 import executorch
 
@@ -15,7 +18,11 @@
 import torchvision.models as torchvision_models
 from executorch.backends.openvino.partitioner import OpenvinoPartitioner
 from executorch.backends.openvino.quantizer import quantize_model
-from executorch.exir import EdgeProgramManager, to_edge_transform_and_lower
+from executorch.exir import (
+    EdgeProgramManager,
+    ExecutorchProgramManager,
+    to_edge_transform_and_lower,
+)
 from executorch.exir.backend.backend_details import CompileSpec
 from executorch.runtime import Runtime
 from sklearn.metrics import accuracy_score
@@ -102,7 +109,7 @@ def load_calibration_dataset(
 
 
 def infer_model(
-    exec_prog: EdgeProgramManager,
+    exec_prog: ExecutorchProgramManager,
     inputs,
     num_iter: int,
     warmup_iter: int,
@@ -111,7 +118,7 @@ def infer_model(
     """
     Executes inference and reports the average timing.
 
-    :param exec_prog: EdgeProgramManager of the lowered model
+    :param exec_prog: ExecutorchProgramManager of the lowered model
     :param inputs: The inputs for the model.
     :param num_iter: The number of iterations to execute inference for timing.
     :param warmup_iter: The number of iterations to execute inference for warmup before timing.
@@ -122,8 +129,11 @@ def infer_model(
     runtime = Runtime.get()
     program = runtime.load_program(exec_prog.buffer)
     method = program.load_method("forward")
+    if method is None:
+        raise ValueError("Load method failed")
 
     # Execute warmup
+    out = None
     for _i in range(warmup_iter):
         out = method.execute(inputs)
 
@@ -137,34 +147,38 @@ def infer_model(
 
     # Save output tensor as raw tensor file
     if output_path:
+        assert out is not None
         torch.save(out, output_path)
 
     # Return average inference timing
     return time_total / float(num_iter)
 
 
 def validate_model(
-    exec_prog: EdgeProgramManager, calibration_dataset: torch.utils.data.DataLoader
+    exec_prog: ExecutorchProgramManager,
+    calibration_dataset: torch.utils.data.DataLoader,
 ) -> float:
     """
     Validates the model using the calibration dataset.
 
-    :param exec_prog: EdgeProgramManager of the lowered model
+    :param exec_prog: ExecutorchProgramManager of the lowered model
     :param calibration_dataset: A DataLoader containing calibration data.
     :return: The accuracy score of the model.
     """
     # Load model from buffer
     runtime = Runtime.get()
     program = runtime.load_program(exec_prog.buffer)
     method = program.load_method("forward")
+    if method is None:
+        raise ValueError("Load method failed")
 
     # Iterate over the dataset and run the executor
-    predictions = []
+    predictions: List[int] = []
     targets = []
     for _idx, data in enumerate(calibration_dataset):
         feature, target = data
         targets.extend(target)
-        out = method.execute((feature,))
+        out = list(method.execute((feature,)))
         predictions.extend(torch.stack(out).reshape(-1, 1000).argmax(-1))
 
     # Check accuracy
@@ -213,12 +227,18 @@ def main(  # noqa: C901
     model = load_model(suite, model_name)
     model = model.eval()
 
+    calibration_dataset: Optional[torch.utils.data.DataLoader] = None
+
     if dataset_path:
         calibration_dataset = load_calibration_dataset(
             dataset_path, batch_size, suite, model, model_name
         )
-        input_shape = tuple(next(iter(calibration_dataset))[0].shape)
-        print(f"Input shape retrieved from the model config: {input_shape}")
+        if calibration_dataset is not None:
+            input_shape = tuple(next(iter(calibration_dataset))[0].shape)
+            print(f"Input shape retrieved from the model config: {input_shape}")
+        else:
+            msg = "Quantization requires a valid calibration dataset"
+            raise ValueError(msg)
     # Ensure input_shape is a tuple
     elif isinstance(input_shape, (list, tuple)):
         input_shape = tuple(input_shape)
@@ -240,7 +260,7 @@ def main(  # noqa: C901
     # Export the model to the aten dialect
     aten_dialect: ExportedProgram = export(model, example_args)
 
-    if quantize:
+    if quantize and calibration_dataset:
         if suite == "huggingface":
             msg = f"Quantization of {suite} models did not support yet."
             raise ValueError(msg)
@@ -251,20 +271,20 @@ def main(  # noqa: C901
             raise ValueError(msg)
 
         subset_size = 300
-        batch_size = calibration_dataset.batch_size
+        batch_size = calibration_dataset.batch_size or 1
         subset_size = (subset_size // batch_size) + int(subset_size % batch_size > 0)
 
         def transform_fn(x):
             return x[0]
 
         quantized_model = quantize_model(
-            aten_dialect.module(),
+            cast(torch.fx.GraphModule, aten_dialect.module()),
             calibration_dataset,
             subset_size=subset_size,
             transform_fn=transform_fn,
         )
 
-        aten_dialect: ExportedProgram = export(quantized_model, example_args)
+        aten_dialect = export(quantized_model, example_args)
 
     # Convert to edge dialect and lower the module to the backend with a custom partitioner
     compile_spec = [CompileSpec("device", device.encode())]
@@ -288,7 +308,7 @@ def transform_fn(x):
             exec_prog.write_to_file(file)
         print(f"Model exported and saved as {model_file_name} on {device}.")
 
-    if validate:
+    if validate and calibration_dataset:
         if suite == "huggingface":
             msg = f"Validation of {suite} models did not support yet."
             raise ValueError(msg)