Rework tensor dispatcher (openvinotoolkit#3306)

### Changes

Introduce tensor_dispatcher instead of functools.singledispatch

- Simplified adding a new function
- Dispatch mechanism based on signature 
- Check names of arguments and annotation 
- Enable mypy for `nncf/tensor`
- Keep signature of function

![image](https://github.com/user-attachments/assets/bc53a27e-49ab-4cd6-816b-07e12e431f60)

### Reason for changes

No signature of functions

![image](https://github.com/user-attachments/assets/7bbb4ec0-4323-49b0-bd2b-dd3ea8fe149b)
Impossible use mypy without cast.

Author: AlexanderDokuchaev

nncf/common/tensor_statistics/statistics.py

@@ -12,7 +12,7 @@
 from abc import ABC
 from abc import abstractmethod
 from collections import Counter
-from typing import Any, Dict, TypeVar, cast
+from typing import Any, Dict, TypeVar
 
 from nncf.tensor import Tensor
 from nncf.tensor import functions as fns
@@ -27,7 +27,7 @@ class TensorStatistic(ABC):
 
     @staticmethod
     def tensor_eq(tensor1: Tensor, tensor2: Tensor, rtol: float = 1e-6) -> bool:
-        return cast(bool, fns.allclose(tensor1, tensor2))
+        return fns.allclose(tensor1, tensor2, rtol=rtol)
 
     @abstractmethod
     def __eq__(self, other: Any) -> bool:

nncf/tensor/README.md

@@ -123,45 +123,35 @@ tensor_a[0:2]  # Tensor(array([[1],[2]]))
 2. Add function to functions module
 
     ```python
-    @functools.singledispatch
-    def foo(a: TTensor, arg1: Type) -> TTensor:
+    @tensor_dispatcher
+    def foo(a: Tensor, arg1: Type) -> Tensor:
         """
         __description__
 
         :param a: The input tensor.
         :param arg1: __description__
         :return: __description__
         """
-        if isinstance(a, tensor.Tensor):
-            return tensor.Tensor(foo(a.data, axis))
-        return NotImplemented(f"Function `foo` is not implemented for {type(a)}")
     ```
 
-    **NOTE** For the case when the first argument has type `List[Tensor]`, use the `_dispatch_list` function. This function dispatches function by first element in the first argument.
-
-    ```python
-    @functools.singledispatch
-    def foo(x: List[Tensor], axis: int = 0) -> Tensor:
-        if isinstance(x, List):
-            unwrapped_x = [i.data for i in x]
-            return Tensor(_dispatch_list(foo, unwrapped_x, axis=axis))
-        raise NotImplementedError(f"Function `foo` is not implemented for {type(x)}")
-    ```
+    **NOTE** The wrapping of the return value depends on the return type annotation of the function.
+    If return type collect `Tensor` than return value will be wrapped in `Tensor` class according annotation,
+    otherwise return value will be returned as is.
 
 3. Add backend specific implementation of method to corresponding module:
 
     - `functions/numpy_*.py`
 
         ```python
-        @_register_numpy_types(fns.foo)
-        def _(a: TType, arg1: Type) -> np.ndarray:
+        @fns.foo.register
+        def _(a: T_NUMPY_ARRAY, arg1: Type) -> T_NUMPY_ARRAY:
             return np.foo(a, arg1)
         ```
 
     - `functions/torch_*.py`
 
         ```python
-        @fns.foo.register(torch.Tensor)
+        @fns.foo.register
         def _(a: torch.Tensor, arg1: Type) -> torch.Tensor:
             return torch.foo(a, arg1)
         ```

nncf/tensor/definitions.py

@@ -12,6 +12,12 @@
 from dataclasses import dataclass
 from enum import Enum
 from enum import auto
+from typing import Optional, Tuple, Union
+
+T_SHAPE_ARRAY = Tuple[int, ...]
+T_SHAPE = Union[int, T_SHAPE_ARRAY]
+T_AXIS = Optional[T_SHAPE]
+T_NUMBER = Union[int, float, bool]
 
 
 class TensorBackend(Enum):
@@ -40,7 +46,7 @@ class TensorDataType(Enum):
     uint4 = auto()
     int4 = auto()
 
-    def is_float(self):
+    def is_float(self) -> bool:
         """
         :return: True if the tensor data type is a floating-point type, else False.
         """

nncf/tensor/functions/__init__.py

@@ -67,7 +67,7 @@
 from nncf.tensor.functions.numeric import zeros_like as zeros_like
 
 
-def _initialize_backends():
+def _initialize_backends() -> None:
     import contextlib
 
     import nncf.tensor.functions.numpy_io