Tensor for PTQ

nncf/experimental/common/tensor_statistics/statistical_functions.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2023 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nncf.experimental.tensor import Tensor
+from nncf.experimental.tensor import functions as fns
+
+
+def mean_per_channel(x: Tensor, axis: int) -> Tensor:
+    """
+    Computes the mean of elements across given channel dimension of Tensor.
+
+    :param x: Tensor to reduce.
+    :param axis: The channel dimensions to reduce.
+    :return: Reduced Tensor.
+    """
+    if len(x.shape) < 3:
+        return fns.mean(x, axis=0)
+    pos_axis = axis + x.ndim if axis < 0 else axis
+    if pos_axis < 0 or pos_axis >= x.ndim:
+        raise ValueError(f"axis {axis} is out of bounds for array of dimension {x.ndim}")
+    axis = tuple(i for i in range(x.ndim) if i != pos_axis)
+    return fns.mean(x, axis=axis)

nncf/experimental/tensor/README.md

@@ -6,7 +6,7 @@ making them more portable and reusable.
 
 ## Usage
 
-The main idea is common algorithms should use wrapped tensors and provide to backend-specific function unwrapped tensor.
+Common algorithms should use wrapped tensors and provide the unwrapped tensor to the backend-specific function.
 
 ### Initialization Tensor
 
@@ -32,6 +32,8 @@ tenor_b = Tensor(np.array([1,2]))
 tensor_a + tenor_b  # Tensor(array([2, 4]))
 ```
 
+**NOTE** Division operations for the numpy backend are performed with warnings disabled for the same for all backends.
+
 ### Comparison operators
 
 All math operations are overrided to operated with wrapped object and return `Tensor`
@@ -55,16 +57,16 @@ nncf_tensor.max()  # Tensor(2)
 All available functions you can found in [functions.py](functions.py).
 
 ```python
-from nncf.experimental.tensor import functions
-functions.max(nncf_tensor)  # Tensor(2)
+from nncf.experimental.tensor import functions as fns
+fns.max(nncf_tensor)  # Tensor(2)
 ```
 
 **NOTE** A function requires at least one positional argument, which is used to dispatch the function
 to the appropriate implementation depending on the type of argument.
 
 ```python
-functions.max(nncf_tensor)  # Correct
-functions.max(a=nncf_tensor)  # TypeError: wrapper requires at least 1 positional argument
+fns.max(nncf_tensor)  # Correct
+fns.max(a=nncf_tensor)  # TypeError: wrapper requires at least 1 positional argument
 ```
 
 ### Loop over Tensor
@@ -100,7 +102,7 @@ tensor_a[0:2]  # Tensor(array([[1],[2]]))
     class Tensor:
         ...
         def foo(self, arg1: Type) -> "Tensor":
-            return functions.foo(self, arg1)
+            return fns.foo(self, arg1)
     ```
 
 2. Add function to [function.py](function.py)
@@ -120,28 +122,36 @@ tensor_a[0:2]  # Tensor(array([[1],[2]]))
         return NotImplemented(f"Function `foo` is not implemented for {type(a)}")
     ```
 
-3. Add function name to `__all__` in [function.py](function.py)
+    **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)}")
+    ```
 
-4. Add backend specific implementation of method to:
+3. Add backend specific implementation of method to:
 
-    - [numpy_function.py](numpy_function.py)
+    - [numpy_function.py](numpy_functions.py)
 
         ```python
-        @functions.foo.register(np.ndarray)
-        @functions.foo.register(np.number)
+        @_register_numpy_types(fns.foo)
         def _(a: TType, arg1: Type) -> np.ndarray:
             return np.foo(a, arg1)
         ```
 
-    - [torch_function.py](torch_function.py)
+    - [torch_function.py](torch_functions.py)
 
         ```python
-        @functions.foo.register(torch.Tensor)
+        @fns.foo.register(torch.Tensor)
         def _(a: torch.Tensor, arg1: Type) -> torch.Tensor:
             return torch.foo(a, arg1)
         ```
 
-5. Add test of method to [test template](tests/shared/test_templates/template_test_nncf_tensor.py) for Tensor class
+4. Add test of method to [test template](../../../tests/shared/test_templates/template_test_nncf_tensor.py) for Tensor class
 
 ### Add new backend