[mypy] Enable mypy for some files in OV and ONNX backends

.github/workflows/mypy.yml

@@ -24,7 +24,7 @@ jobs:
           python-version: "3.12"
       - name: Install NNCF
         run: |
-          pip install . torch -c constraints.txt
+          pip install . torch onnx onnxruntime openvino -c constraints.txt
       - name: Install mypy
         run: pip install mypy==1.8.0
       - name: Run mypy

pyproject.toml

@@ -74,10 +74,7 @@ strict = true
 files = ["src"]
 exclude = [
     "src/nncf/experimental/*",
-    "src/nncf/onnx/engine.py",
-    "src/nncf/onnx/graph/metatypes/groups.py",
     "src/nncf/onnx/graph/metatypes/onnx_metatypes.py",
-    "src/nncf/onnx/graph/model_metadata.py",
     "src/nncf/onnx/graph/model_transformer.py",
     "src/nncf/onnx/graph/model_utils.py",
     "src/nncf/onnx/graph/nncf_graph_builder.py",
@@ -86,37 +83,25 @@ exclude = [
     "src/nncf/onnx/graph/passes.py",
     "src/nncf/onnx/graph/transformations/command_creation.py",
     "src/nncf/onnx/graph/transformations/commands.py",
-    "src/nncf/onnx/hardware/config.py",
     "src/nncf/onnx/quantization/backend_parameters.py",
-    "src/nncf/onnx/quantization/default_quantization.py",
     "src/nncf/onnx/quantization/quantize_model.py",
     "src/nncf/onnx/quantization/quantizer_parameters.py",
     "src/nncf/onnx/statistics/aggregator.py",
-    "src/nncf/openvino/cpu_info.py",
-    "src/nncf/openvino/engine.py",
-    "src/nncf/openvino/graph/layer_attributes.py",
     "src/nncf/openvino/graph/layout.py",
-    "src/nncf/openvino/graph/metatypes/groups.py",
     "src/nncf/openvino/graph/metatypes/openvino_metatypes.py",
-    "src/nncf/openvino/graph/model_builder.py",
     "src/nncf/openvino/graph/model_transformer.py",
-    "src/nncf/openvino/graph/model_utils.py",
     "src/nncf/openvino/graph/nncf_graph_builder.py",
     "src/nncf/openvino/graph/node_utils.py",
     "src/nncf/openvino/graph/transformations/command_creation.py",
-    "src/nncf/openvino/graph/transformations/commands.py",
-    "src/nncf/openvino/hardware/config.py",
     "src/nncf/openvino/optimized_functions/functions.py",
     "src/nncf/openvino/optimized_functions/models.py",
     "src/nncf/openvino/quantization/backend_parameters.py",
     "src/nncf/openvino/quantization/default_quantization.py",
     "src/nncf/openvino/quantization/quantize_ifmodel.py",
     "src/nncf/openvino/quantization/quantize_model.py",
-    "src/nncf/openvino/rt_info.py",
     "src/nncf/openvino/statistics/aggregator.py",
     "src/nncf/openvino/statistics/builders.py",
     "src/nncf/openvino/statistics/collectors.py",
-    "src/nncf/pruning/prune_model.py",
     "src/nncf/quantization/algorithms/accuracy_control/algorithm.py",
     "src/nncf/quantization/algorithms/accuracy_control/backend.py",
     "src/nncf/quantization/algorithms/accuracy_control/evaluator.py",

src/nncf/common/factory.py

@@ -39,7 +39,7 @@ def build_graph(model: TModel, *, example_input: Any = None) -> NNCFGraph:
     """
     model_backend = get_backend(model)
     if model_backend == BackendType.ONNX:
-        from onnx import ModelProto  # type: ignore
+        from onnx import ModelProto
 
         from nncf.onnx.graph.nncf_graph_builder import GraphConverter as ONNXGraphConverter
 

src/nncf/common/graph/transformations/commands.py

@@ -139,7 +139,7 @@ def __init__(self, target_type: TargetType):
     def type(self) -> TargetType:
         return self._target_type
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         return isinstance(other, TargetPoint) and self.type == other.type
 
     def __str__(self) -> str:

src/nncf/common/tensor_statistics/aggregator.py

@@ -47,7 +47,7 @@ class StatisticsAggregator(ABC):
 
     def __init__(self, dataset: Dataset):
         self.dataset = dataset
-        self.stat_subset_size = None
+        self.stat_subset_size: int | None = None
         self.statistic_points = StatisticPointsContainer()
 
     def _get_iterations_number(self) -> int | None:

src/nncf/common/tensor_statistics/statistic_point.py

@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from collections import UserDict
-from typing import Any, Callable, Generator, cast
+from typing import Callable, Generator
 
 from nncf.common.graph.transformations.commands import TargetPoint
 from nncf.common.tensor_statistics.collectors import TensorCollector
@@ -28,15 +28,17 @@ def __init__(self, target_point: TargetPoint, tensor_collector: TensorCollector,
         self.target_point = target_point
         self.algorithm_to_tensor_collectors = {algorithm: [tensor_collector]}
 
-    def __eq__(self, other: Any) -> bool:
-        return cast(
-            bool,
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, StatisticPoint):
+            raise NotImplementedError
+
+        return (
             self.target_point == other.target_point
-            and self.algorithm_to_tensor_collectors == other.self.algorithm_to_tensor_collectors,
+            and self.algorithm_to_tensor_collectors == other.algorithm_to_tensor_collectors
         )
 
 
-class StatisticPointsContainer(UserDict):  # type: ignore
+class StatisticPointsContainer(UserDict[str, list[StatisticPoint]]):
     """
     Container with iteration interface for handling a composition of StatisticPoint.
     """

src/nncf/common/tensor_statistics/statistics.py

@@ -83,7 +83,7 @@ class MinMaxTensorStatistic(TensorStatistic):
     min_values: Tensor
     max_values: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, MinMaxTensorStatistic):
             return fns.allclose(self.min_values, other.min_values) and fns.allclose(self.max_values, other.max_values)
         return False
@@ -103,7 +103,7 @@ class AbsMaxTensorStatistic(TensorStatistic):
 
     abs_max: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, AbsMaxTensorStatistic):
             return fns.allclose(self.abs_max, other.abs_max)
         return False
@@ -121,7 +121,7 @@ def __init__(self, mean_values: Tensor, shape: Tensor) -> None:
         self.mean_values = mean_values
         self.shape = tuple(shape.tolist())
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, MeanTensorStatistic):
             return self.shape == other.shape and fns.allclose(self.mean_values, other.mean_values)
         return False
@@ -147,7 +147,7 @@ class MedianMADTensorStatistic(TensorStatistic):
     median_values: Tensor
     mad_values: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, MedianMADTensorStatistic):
             return fns.allclose(self.median_values, other.median_values) and fns.allclose(
                 self.mad_values, other.mad_values
@@ -168,7 +168,7 @@ class PercentileTensorStatistic(TensorStatistic):
 
     percentile_vs_values_dict: dict[int, Tensor]
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, PercentileTensorStatistic):
             if Counter(self.percentile_vs_values_dict.keys()) != Counter(other.percentile_vs_values_dict.keys()):
                 return False
@@ -202,7 +202,7 @@ class RawTensorStatistic(TensorStatistic):
 
     values: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, RawTensorStatistic):
             return fns.allclose(self.values, other.values)
         return False
@@ -214,7 +214,7 @@ class HessianTensorStatistic(TensorStatistic):
 
     hessian: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, HessianTensorStatistic):
             return fns.allclose(self.hessian, other.hessian)
         return False
@@ -226,7 +226,7 @@ class MeanVarianceTensorStatistic(TensorStatistic):
 
     mean_variance: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, MeanVarianceTensorStatistic):
             return fns.allclose(self.mean_variance, other.mean_variance)
         return False
@@ -238,7 +238,7 @@ class MaxVarianceTensorStatistic(TensorStatistic):
 
     max_variance: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, MaxVarianceTensorStatistic):
             return fns.allclose(self.max_variance, other.max_variance)
         return False
@@ -250,7 +250,7 @@ class MeanMagnitudeTensorStatistic(TensorStatistic):
 
     mean_magnitude: Tensor
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if isinstance(other, MeanMagnitudeTensorStatistic):
             return fns.allclose(self.mean_magnitude, other.mean_magnitude)
         return False
@@ -264,7 +264,7 @@ class WCTensorStatistic(TensorStatistic):
     mean_values: list[Tensor]
     shape_values: list[tuple[Tensor, ...]]
 
-    def __eq__(self, other: Any) -> bool:
+    def __eq__(self, other: object) -> bool:
         if not isinstance(other, WCTensorStatistic):
             return False
         shapes_equal = all(self.shape_values[i] == other.shape_values[i] for i in range(len(self.mean_values)))

src/nncf/common/utils/backend.py

@@ -80,7 +80,7 @@ def is_onnx_model(model: Any) -> bool:
     :param model: A target model.
     :return: True if the model is an instance of onnx.ModelProto, otherwise False.
     """
-    import onnx  # type: ignore
+    import onnx
 
     return isinstance(model, onnx.ModelProto)
 

src/nncf/onnx/engine.py

@@ -42,7 +42,7 @@ def __init__(self, model: ModelProto, **rt_session_options: Any):
         for inp in self.sess.get_inputs():
             self.input_names.add(inp.name)
 
-    def infer(self, input_data: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
+    def infer(self, input_data: dict[str, np.ndarray[Any, Any]]) -> dict[str, np.ndarray[Any, Any]]:
         """
         Runs model on the provided input via ONNXRuntime InferenceSession.
         Returns the dictionary of model outputs by node names.

src/nncf/onnx/graph/model_transformer.py

@@ -10,6 +10,7 @@
 # limitations under the License.
 from collections import Counter
 from copy import deepcopy
+from typing import Any
 
 import numpy as np
 import onnx
@@ -37,7 +38,7 @@
 from nncf.onnx.graph.transformations.commands import ONNXQuantizerInsertionCommand
 
 
-class ONNXModelTransformer(ModelTransformer):
+class ONNXModelTransformer(ModelTransformer[onnx.ModelProto]):
     """
     Applies transformations upon ONNX model.
     ModelTransformer should be created once for a particular model,
@@ -546,7 +547,7 @@ def _apply_multiply_insertion_transformations(
         return model
 
 
-def set_initializer(initializer_name: str, model: onnx.ModelProto, new_value: np.ndarray) -> None:
+def set_initializer(initializer_name: str, model: onnx.ModelProto, new_value: np.ndarray[Any, Any]) -> None:
     """
     Updates the initializer tensor in the ONNX model.
     :param initializer_name: Name of the initializer tensor to update.

src/nncf/onnx/graph/onnx_helper.py

@@ -46,7 +46,7 @@ def get_name_to_node_map(model: onnx.ModelProto) -> dict[str, onnx.NodeProto]:
     return {node.name: node for node in model.graph.node}
 
 
-def get_edge_info_mapping(model: onnx.ModelProto) -> dict[str, onnx.ValueInfoProto]:
+def get_edge_info_mapping(model: onnx.ModelProto) -> dict[str, onnx.ValueInfoProto | onnx.TensorProto]:
     """
     Returns mapping from edge name to the edge info.
 
@@ -195,7 +195,7 @@ def get_tensor(model: onnx.ModelProto, tensor_name: str) -> onnx.TensorProto:
     raise nncf.ValidationError(msg)
 
 
-def get_tensor_value(model: onnx.ModelProto, tensor_name: str) -> np.ndarray:
+def get_tensor_value(model: onnx.ModelProto, tensor_name: str) -> np.ndarray[Any, Any]:
     """
     Returns tensor value of a tensor with the name 'tensor_name'.
 
@@ -207,7 +207,7 @@ def get_tensor_value(model: onnx.ModelProto, tensor_name: str) -> np.ndarray:
     return get_array_from_tensor(model, tensor)
 
 
-def get_array_from_tensor(model: onnx.ModelProto, tensor: onnx.TensorProto) -> np.ndarray:
+def get_array_from_tensor(model: onnx.ModelProto, tensor: onnx.TensorProto) -> np.ndarray[Any, Any]:
     """
     Returns the data from an ONNX tensor as NumPy array.
 
@@ -230,7 +230,7 @@ def get_edge_shape(edge: onnx.ValueInfoProto | onnx.TensorProto) -> list[int]:
     if isinstance(edge, onnx.TensorProto):
         return list(edge.dims)
     tensor_type = edge.type.tensor_type
-    shape = []
+    shape: list[int] = []
     if not tensor_type.HasField("shape"):
         return shape  # shape is unknown
 
@@ -310,7 +310,7 @@ def is_node_has_shared_weight(
     return len(nodes) > 1
 
 
-def pack_4_bits(tensor: np.ndarray) -> np.ndarray:
+def pack_4_bits(tensor: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
     """
     Apply packing based on the rule - https://onnx.ai/onnx/technical/int4.html#packing-and-unpacking
     :param tensor: Tensor to pack.
@@ -333,7 +333,7 @@ def pack_4_bits(tensor: np.ndarray) -> np.ndarray:
     return packed_tensor
 
 
-def pack_int4_to_uint8(weight: np.ndarray, block_size: int, signed: bool) -> np.ndarray:
+def pack_int4_to_uint8(weight: np.ndarray[Any, Any], block_size: int, signed: bool) -> np.ndarray[Any, Any]:
     """
     Returns `weight` that is stored as uint8 with shape (N, n_blocks_per_col, blob_size) in which:
         - n_blocks_per_col = CeilDiv(K, block_size)

src/nncf/onnx/graph/transformations/command_creation.py

@@ -9,6 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Any
+
 import numpy as np
 
 from nncf.common.graph.graph import NNCFGraph
@@ -21,7 +23,7 @@
 from nncf.onnx.graph.transformations.commands import ONNXTargetPoint
 
 
-def create_bias_correction_command(node: NNCFNode, bias_value: np.ndarray) -> ONNXInitializerUpdateCommand:
+def create_bias_correction_command(node: NNCFNode, bias_value: np.ndarray[Any, Any]) -> ONNXInitializerUpdateCommand:
     """
      Creates bias correction command.
 
@@ -51,13 +53,13 @@ def create_command_to_remove_quantizer(quantizer_node: NNCFNode) -> ONNXQDQNodeR
 
     @staticmethod
     def create_command_to_update_bias(
-        node_with_bias: NNCFNode, bias_value: np.ndarray, nncf_graph: NNCFGraph
+        node_with_bias: NNCFNode, bias_value: np.ndarray[Any, Any], nncf_graph: NNCFGraph
     ) -> ONNXInitializerUpdateCommand:
         return create_bias_correction_command(node_with_bias, bias_value)
 
     @staticmethod
     def create_command_to_update_weight(
-        node_with_weight: NNCFNode, weight_value: np.ndarray, weight_port_id: int
+        node_with_weight: NNCFNode, weight_value: np.ndarray[Any, Any], weight_port_id: int
     ) -> ONNXInitializerUpdateCommand:
         target_point = ONNXTargetPoint(TargetType.LAYER, node_with_weight.node_name, weight_port_id)
         return ONNXInitializerUpdateCommand(target_point, weight_value)
@@ -71,7 +73,7 @@ def multiply_insertion_command(
         source_node: NNCFNode,
         destination_nodes: list[NNCFNode],
         source_out_port: int,
-        scale_value: np.ndarray,
+        scale_value: np.ndarray[Any, Any],
         multiply_node_name: str,
     ) -> ONNXMultiplyInsertionCommand:
         target_point = ONNXTargetPoint(TargetType.POST_LAYER_OPERATION, source_node.node_name, source_out_port)

src/nncf/onnx/graph/transformations/commands.py

@@ -10,6 +10,8 @@
 # limitations under the License.
 
 
+from typing import Any
+
 import numpy as np
 
 from nncf.common.graph.transformations.commands import Command
@@ -34,7 +36,7 @@ def __init__(self, target_type: TargetType, target_node_name: str, port_id: int
         self.target_node_name = target_node_name
         self.port_id = port_id
 
-    def __eq__(self, other: "ONNXTargetPoint") -> bool:
+    def __eq__(self, other: object) -> bool:
         return (
             isinstance(other, ONNXTargetPoint)
             and self.type == other.type
@@ -76,7 +78,7 @@ class ONNXInitializerUpdateCommand(TransformationCommand):
     Update initializer in the value.
     """
 
-    def __init__(self, target_point: ONNXTargetPoint, new_value: np.ndarray):
+    def __init__(self, target_point: ONNXTargetPoint, new_value: np.ndarray[Any, Any]):
         """
         :param target_point: Target point.
         :param new_value: New value for initializer.