[mypy] several files from src/nncf/torch (#3866)

### Changes

Add file to mypy check: 
```
"src/custom_version.py",
"src/nncf/torch/engine.py",
"src/nncf/torch/functions.py",
"src/nncf/torch/model_creation.py",
"src/nncf/torch/node_utils.py",
"src/nncf/torch/strip.py",
"src/nncf/torch/utils.py",
```
Remove unused function:
```
BaseQuantizer.apply_minmax_init
get_flat_tensor_contents_string
```
Moved `sumlike` from `nncf.torch.utils` to
`tests.torch.quantization.reference` as it is used only here

Author: AlexanderDokuchaev

pyproject.toml

@@ -76,10 +76,17 @@ strict = true
 # https://github.com/hauntsaninja/no_implicit_optional
 implicit_optional = true
 files = [
+    "src/custom_version.py",
     "src/nncf/api",
     "src/nncf/data",
     "src/nncf/common",
     "src/nncf/torch/function_hook",
+    "src/nncf/torch/engine.py",
+    "src/nncf/torch/functions.py",
+    "src/nncf/torch/model_creation.py",
+    "src/nncf/torch/node_utils.py",
+    "src/nncf/torch/strip.py",
+    "src/nncf/torch/utils.py",
     "src/nncf/quantization/*py",
     "src/nncf/telemetry/",
     "src/nncf/tensor/",

src/custom_version.py

@@ -104,7 +104,11 @@ def __getattr__(name: str) -> str:
         # Rewrite version.py to pass custom version to package
         if os.environ.get("_PYPROJECT_HOOKS_BUILD_BACKEND"):
             content = Path(NNCF_VERSION_FILE).read_text()
-            version_str = re.search(r"^__version__ = ['\"][^'\"]*['\"]", content, re.M).group(0)
+            version_match = re.search(r"^__version__ = ['\"][^'\"]*['\"]", content, re.M)
+            if version_match is None:
+                msg = "Unable to find version string."
+                raise RuntimeError(msg)
+            version_str = version_match.group(0)
             content = content.replace(version_str, f'__version__ = "{version}"')
             Path(NNCF_VERSION_FILE).write_text(content)
 

src/nncf/torch/engine.py

@@ -9,9 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Union
+from typing import Any
 
-import torch
 from torch import nn
 
 from nncf.common.engine import Engine
@@ -34,9 +33,7 @@ def __init__(self, model: nn.Module):
         if get_backend(model) == BackendType.TORCH:
             self._model.eval()
 
-    def infer(
-        self, input_data: Union[torch.Tensor, tuple[torch.Tensor], dict[str, torch.Tensor]]
-    ) -> Union[torch.Tensor, dict[str, Any]]:
+    def infer(self, input_data: Any) -> Any:
         """
         Runs Torch model on the provided input.
 

src/nncf/torch/functions.py

@@ -13,7 +13,7 @@
 import torch
 
 
-def clamp(x, low, high):
+def clamp(x: torch.Tensor, low: torch.Tensor, high: torch.Tensor) -> torch.Tensor:
     return torch.max(torch.min(x, high), low)
 
 
@@ -28,5 +28,5 @@ def forward(ctx: Any, input_: torch.Tensor, threshold: float = 0.5) -> torch.Ten
         return output
 
     @staticmethod
-    def backward(ctx: Any, *grad_outputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, None]:
+    def backward(ctx: Any, *grad_outputs: torch.Tensor) -> tuple[torch.Tensor, None]:
         return grad_outputs[0], None

src/nncf/torch/node_utils.py

@@ -29,6 +29,9 @@ def get_activation_channel_axis(node: NNCFNode, port_id: int) -> int:
         raise nncf.InternalError(msg)
 
     if node.metatype not in [PTMatMulMetatype, PTAddmmMetatype]:
+        if not isinstance(node.metatype.output_channel_axis, int):
+            msg = f"Node metatype {node.metatype} does not have defined output channel axis"
+            raise nncf.InternalError(msg)
         return node.metatype.output_channel_axis
 
     if port_id == 0:
@@ -38,5 +41,5 @@ def get_activation_channel_axis(node: NNCFNode, port_id: int) -> int:
         # W(port:0) * X(port:1): [... , C_OUT, C_IN] * [... , C_IN, ...]
         return -2
 
-    msg = f"Port id for a {node.metatype} operation is expected to be in [0, 1], {port_id} recieved"
+    msg = f"Port id for a {node.metatype} operation is expected to be in [0, 1], {port_id} received"
     raise nncf.InternalError(msg)

src/nncf/torch/quantization/layers.py

@@ -23,7 +23,6 @@
 
 import nncf
 from nncf.common.graph import NNCFNodeName
-from nncf.common.logging import nncf_logger
 from nncf.common.quantization.quantizer_setup import QuantizationPointId
 from nncf.common.quantization.quantizer_setup import QuantizerSetupBase
 from nncf.common.quantization.quantizers import calculate_asymmetric_level_ranges
@@ -34,7 +33,6 @@
 from nncf.common.quantization.structs import QuantizerSpec
 from nncf.common.utils.debug import is_debug
 from nncf.common.utils.registry import Registry
-from nncf.torch.functions import clamp
 from nncf.torch.graph.transformations.commands import PTTargetPoint
 from nncf.torch.graph.transformations.commands import TargetType
 from nncf.torch.layer_utils import COMPRESSION_MODULES
@@ -56,8 +54,6 @@
 from nncf.torch.quantization.quantize_functions import unpack_uint4
 from nncf.torch.return_types import maybe_get_values_from_torch_return_type
 from nncf.torch.return_types import maybe_wrap_to_torch_return_type
-from nncf.torch.utils import get_flat_tensor_contents_string
-from nncf.torch.utils import get_model_device
 from nncf.torch.utils import is_tracing_state
 from nncf.torch.utils import no_jit_trace
 
@@ -464,29 +460,6 @@ def reset_call_counter(self):
     def get_trainable_params(self) -> dict[str, torch.Tensor]:
         return {}
 
-    def apply_minmax_init(self, min_values: torch.Tensor, max_values: torch.Tensor, log_module_name: str = None):
-        """min_values and max_values must have the same shape as specified in self.scale_shape"""
-        if self.initialized:
-            nncf_logger.debug(f"Skipped initializing {log_module_name} - loaded from checkpoint")
-            return
-
-        if torch.all(torch.isinf(min_values)) or torch.all(torch.isinf(max_values)):
-            msg = f"Statistics are not collected for {log_module_name}"
-            raise ValueError(msg)
-
-        if torch.any(torch.eq(min_values, np.inf)) or torch.any(torch.eq(max_values, -np.inf)):
-            msg = f"Some of the values in statistics have infinite value for {log_module_name}"
-            raise ValueError(msg)
-
-        own_device = get_model_device(self)
-        min_values = min_values.to(own_device)
-        max_values = max_values.to(own_device)
-        self._apply_minmax_init(min_values, max_values, log_module_name)
-
-    @abstractmethod
-    def _apply_minmax_init(self, min_values: torch.Tensor, max_values: torch.Tensor, log_module_name: str = None):
-        pass
-
     @abstractmethod
     def set_levels(self):
         """
@@ -795,26 +768,6 @@ def quantize(self, x, execute_traced_op_as_identity: bool = False):
     def get_trainable_params(self) -> dict[str, torch.Tensor]:
         return {self.SCALE_PARAM_NAME: self.scale}
 
-    def _apply_minmax_init(self, min_values, max_values, log_module_name: str = None):
-        sign = torch.any(torch.lt(min_values, 0))
-        if self._signedness_to_force is not None and sign != self._signedness_to_force:
-            nncf_logger.debug(f"Forcing signed to {self._signedness_to_force} for module {log_module_name}")
-            sign = self._signedness_to_force
-        self.signed = sign
-
-        abs_max = torch.max(torch.abs(max_values), torch.abs(min_values))
-        SCALE_LOWER_THRESHOLD = 0.1
-        mask = torch.gt(abs_max, SCALE_LOWER_THRESHOLD)
-        self._scale_param_storage.data = torch.where(
-            mask, abs_max, SCALE_LOWER_THRESHOLD * torch.ones_like(self._scale_param_storage)
-        )
-        if self._is_using_log_scale_storage:
-            self._scale_param_storage.data.log_()
-
-        nncf_logger.debug(
-            f"Set sign: {self.signed} and scale: {get_flat_tensor_contents_string(self.scale)} for {log_module_name}"
-        )
-
     def broadcast_initialized_params(self, src: int = 0):
         super().broadcast_initialized_params(src)
         distributed.broadcast(self._scale_param_storage, src=src)
@@ -996,22 +949,6 @@ def get_trainable_params(self) -> dict[str, torch.Tensor]:
             self.INPUT_RANGE_PARAM_NAME: self.input_range,
         }
 
-    def _apply_minmax_init(self, min_values, max_values, log_module_name: str = None):
-        ranges = max_values - min_values
-        max_range = torch.max(max_values - min_values)
-        eps = 1e-2
-        correction = (clamp(ranges, low=eps * max_range, high=max_range) - ranges) * 0.5
-        self._input_range_param_storage.data = (ranges + 2 * correction).data
-        if self._is_using_log_scale_storage:
-            self._input_range_param_storage.data.log_()
-
-        self.input_low.data = (min_values - correction).data
-
-        nncf_logger.debug(
-            f"Set input_low: {get_flat_tensor_contents_string(self.input_low)} "
-            f"and input_range: {get_flat_tensor_contents_string(self.input_range)} for {log_module_name}"
-        )
-
     def broadcast_initialized_params(self, src: int = 0):
         super().broadcast_initialized_params(src)
         distributed.broadcast(self.input_low, src)

src/nncf/torch/quantization/reference.py

@@ -17,11 +17,38 @@
 
 import nncf
 from nncf.torch.utils import CompilationWrapper
-from nncf.torch.utils import sum_like
 
 GeneralizedTensor = TypeVar("GeneralizedTensor", torch.Tensor, np.ndarray)
 
 
+def fp32_accum_wrapper(func):
+    def wrapper(tensor_to_sum, ret_tensor):
+        half = tensor_to_sum.dtype == np.float16
+        if half:
+            tensor_to_sum = tensor_to_sum.astype(np.float32)
+        retval = func(tensor_to_sum, ret_tensor)
+        if half:
+            retval = retval.astype(np.float16)
+        return retval
+
+    return wrapper
+
+
+@fp32_accum_wrapper
+def sum_like(tensor_to_sum, ref_tensor):
+    """Warning: may modify tensor_to_sum"""
+    if ref_tensor.size == 1:
+        return tensor_to_sum.sum()
+
+    for dim, size in enumerate(ref_tensor.shape):
+        if size == 1:
+            if isinstance(tensor_to_sum, np.ndarray):
+                tensor_to_sum = tensor_to_sum.sum(dim, keepdims=True)
+            else:
+                tensor_to_sum = tensor_to_sum.sum(dim, keepdim=True)
+    return tensor_to_sum
+
+
 class ReferenceBackendType(Enum):
     NUMPY = "numpy"
     TORCH = "torch"

src/nncf/torch/utils.py

@@ -11,7 +11,6 @@
 from contextlib import contextmanager
 from typing import Any, Callable, Generator
 
-import numpy as np
 import torch
 from torch.nn import Module
 
@@ -20,57 +19,33 @@
 from nncf.common.utils.os import is_windows
 
 
-def is_tracing_state():
-    return torch._C._get_tracing_state() is not None
-
-
-class no_jit_trace:
-    def __enter__(self):
-        self.state = torch._C._get_tracing_state()
-        torch._C._set_tracing_state(None)
-
-    def __exit__(self, *args):
-        torch._C._set_tracing_state(self.state)
-        self.state = None
-
-
-def fp32_accum_wrapper(func):
-    def wrapper(tensor_to_sum, ret_tensor):
-        half = tensor_to_sum.dtype == np.float16
-        if half:
-            tensor_to_sum = tensor_to_sum.astype(np.float32)
-        retval = func(tensor_to_sum, ret_tensor)
-        if half:
-            retval = retval.astype(np.float16)
-        return retval
-
-    return wrapper
-
+def is_tracing_state() -> bool:
+    """
+    Checks whether the current execution context is being traced by torch.jit.
 
-@fp32_accum_wrapper
-def sum_like(tensor_to_sum, ref_tensor):
-    """Warning: may modify tensor_to_sum"""
-    if ref_tensor.size == 1:
-        return tensor_to_sum.sum()
+    :return: True if the current thread is being traced, False otherwise.
+    """
+    return torch._C._get_tracing_state() is not None
 
-    for dim, size in enumerate(ref_tensor.shape):
-        if size == 1:
-            if isinstance(tensor_to_sum, np.ndarray):
-                tensor_to_sum = tensor_to_sum.sum(dim, keepdims=True)
-            else:
-                tensor_to_sum = tensor_to_sum.sum(dim, keepdim=True)
-    return tensor_to_sum
 
+@contextmanager
+def no_jit_trace() -> Generator[None, None, None]:
+    """
+    Context manager and decorator to temporarily disable PyTorch JIT tracing.
 
-def get_flat_tensor_contents_string(input_tensor):
-    retval = "["
-    for idx, el in enumerate(input_tensor.view(-1)):
-        if idx >= 10:
-            retval += f"... (first 10/{len(input_tensor.view(-1))} elements shown only) "
-            break
-        retval += f"{el.item():.4f}, "
-    retval += "]"
-    return retval
+    When used, any operations performed within this scope will not be recorded
+    in the TorchScript graph, even if the code is currently being executed
+    via `torch.jit.trace`.
+    """
+    # Capture the original state
+    original_state = torch._C._get_tracing_state()
+    try:
+        # Disable tracing
+        torch._C._set_tracing_state(None)  # type: ignore[attr-defined]
+        yield
+    finally:
+        # Restore state regardless of whether an error occurred
+        torch._C._set_tracing_state(original_state)  # type: ignore[attr-defined]
 
 
 class _ModuleState:
@@ -97,7 +72,7 @@ def save_module_state(module: Module) -> _ModuleState:
     return _ModuleState(module)
 
 
-def load_module_state(base_module: Module, state: _ModuleState, strict=False) -> None:
+def load_module_state(base_module: Module, state: _ModuleState, strict: bool = False) -> None:
     for name, module in base_module.named_modules():
         try:
             module.train(state.training_state[name])
@@ -114,7 +89,7 @@ def load_module_state(base_module: Module, state: _ModuleState, strict=False) ->
 
 
 @contextmanager
-def training_mode_switcher(model: Module, is_training: bool = True):
+def training_mode_switcher(model: Module, is_training: bool = True) -> Generator[None, None, None]:
     saved_state = save_module_state(model)
     model.train(is_training)
     try:

tests/common/test_statistic_collector.py

@@ -150,7 +150,7 @@ def test_duplicated_statistics_are_merged():
     target_inputs = TensorCollector.get_tensor_collector_inputs(outputs, output_info)
     collector.register_inputs(target_inputs)
 
-    # Check aggregators recieved inputs as expected
+    # Check aggregators received inputs as expected
     assert aggregators[0]._collected_samples == 1
     for aggregator in aggregators[1:]:
         assert aggregator._collected_samples == 0
@@ -160,7 +160,7 @@ def test_duplicated_statistics_are_merged():
 
     statistics = collector.get_statistics()
 
-    # Check aggregators recieved correct inputs
+    # Check aggregators received correct inputs
     assert len(statistics) == 6
     for k in "ABC":
         assert statistics[k] == Tensor(np.array(5))