[ONNX][SmoothQuant] Introduce new axes and axes_mode parameters

src/nncf/experimental/common/tensor_statistics/collectors.py

@@ -17,6 +17,7 @@
 from collections import defaultdict
 from collections import deque
 from copy import deepcopy
+from enum import Enum
 from typing import Any, Optional, TypeVar, Union
 
 import nncf
@@ -35,6 +36,50 @@
 
 InplaceInsertionFNType = TypeVar("InplaceInsertionFNType")
 AggregationAxes = tuple[int, ...]
+Axes = tuple[int, ...]
+
+
+class AxesMode(Enum):
+    """
+    Represents different strategies for handling tensor axes.
+
+    :param REDUCTION: Indicates that the specified axes should be reduced during an operation.
+    :param KEEP: Indicates that the specified axes should be preserved and not reduced during
+        an operation.
+    """
+
+    REDUCTION = "reduction"
+    KEEP = "keep"
+
+
+def determine_reduction_axes(
+    ndim: int, axes: Optional[Axes] = None, axes_mode: AxesMode = AxesMode.REDUCTION
+) -> ReductionAxes:
+    """
+    Determines the set of axes along which a reduction operation should be performed
+    based on the specified axes mode.
+
+    :param ndim: The number of dimensions in the input tensor.
+    :param axes: The axes specified for the reduction operation. If `None`, all axes
+        are considered (i.e., `tuple(range(ndim))`).
+
+    :param axes_mode: Defines how the specified axes are interpreted:
+        - `AxesMode.REDUCTION`: the given axes will be reduced.
+        - `AxesMode.KEEP`: all axes except the specified ones will be reduced.
+    :return: The resolved set of axes along which the reduction operation should be performed.
+    """
+    if axes is None:
+        return tuple(range(ndim))
+
+    if axes_mode == AxesMode.REDUCTION:
+        return axes
+
+    all_axes = tuple(range(ndim))
+    if len(all_axes) > 1:
+        # Ensure that all axes have positive values
+        keep_axes = tuple(all_axes[i] for i in axes)
+        return tuple(set(all_axes) - set(keep_axes))
+    return ()
 
 
 class TensorReducerBase(ABC):
@@ -43,13 +88,21 @@ class TensorReducerBase(ABC):
     the specified rule. Could handle tensors inplace or out of place.
     """
 
-    def __init__(self, reduction_axes: Optional[ReductionAxes] = None, inplace: bool = False):
+    def __init__(
+        self,
+        axes: Optional[Axes] = None,
+        axes_mode: AxesMode = AxesMode.REDUCTION,
+        inplace: bool = False,
+    ):
         """
-        :param reduction_axes: Reduction axes for reduction calculation. Equal to list(range(len(input.shape)))
-            if empty.
+        :param axes: The axes along which the reduction operation should be applied.
+            If `None`, the operation will be applied to all axes (i.e., `tuple(range(tensor.ndim))`).
+        :param axes_mode: Determines how the specified `axes` are treated during the operation.
+            Use `AxesMode.REDUCTION` to reduce over the given axes, or `AxesMode.KEEP` to preserve them.
         :param inplace: Whether should be calculated inplace or out of place.
         """
-        self._reduction_axes = reduction_axes
+        self._axes = axes
+        self._axes_mode = axes_mode
         self._inplace = inplace
         self._keepdims = True
 
@@ -97,17 +150,13 @@ def __call__(self, x: list[Tensor]):
     def __eq__(self, __o: object) -> bool:
         return (
             isinstance(__o, self.__class__)
-            and self._reduction_axes == __o._reduction_axes
+            and self._axes == __o._axes
+            and self._axes_mode == __o._axes_mode
             and self._inplace == __o.inplace
         )
 
     def __hash__(self) -> int:
-        return hash((self.__class__.__name__, self.inplace, self._reduction_axes))
-
-    def _get_reduction_axes(self, tensor: Tensor) -> ReductionAxes:
-        if self._reduction_axes is not None:
-            return self._reduction_axes
-        return tuple(range(len(tensor.shape)))
+        return hash((self.__class__.__name__, self.inplace, self._axes, self._axes_mode))
 
 
 class AggregatorBase:
@@ -444,92 +493,94 @@ def get_inplace_fn(self) -> Optional[InplaceInsertionFNType]:
 class MinReducer(TensorReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = x[0]
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         return [fns.min(x, reduction_axes, keepdims=self._keepdims)]
 
 
 class MaxReducer(TensorReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = x[0]
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         return [fns.max(x, reduction_axes, keepdims=self._keepdims)]
 
 
 class AbsMaxReducer(TensorReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = fns.abs(x[0])
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         return [fns.max(x, reduction_axes, keepdims=self._keepdims)]
 
 
 class MeanReducer(TensorReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = x[0]
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         return [fns.mean(x, reduction_axes, keepdims=self._keepdims)]
 
 
 class MeanVarianceReducer(TensorReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = x[0]
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         variance = fns.var(x, reduction_axes)
         return [fns.mean(variance)]
 
 
 class MaxVarianceReducer(TensorReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = x[0]
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         variance = fns.var(x, reduction_axes)
         return [fns.max(variance)]
 
 
 class MeanAbsMaxReducer(TensorReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = fns.abs(x[0])
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         abs_max = fns.max(x, reduction_axes, keepdims=self._keepdims)
         return [fns.mean(abs_max)]
 
 
 class QuantileReducerBase(TensorReducerBase):
     def __init__(
         self,
-        reduction_axes: Optional[ReductionAxes] = None,
+        axes: Optional[Axes] = None,
+        axes_mode: AxesMode = AxesMode.REDUCTION,
         quantile: Optional[Union[float, tuple[float]]] = None,
         inplace: bool = False,
     ):
-        super().__init__(reduction_axes=reduction_axes, inplace=False)
+        super().__init__(axes, axes_mode, False)
         self._quantile = (0.01, 0.99) if quantile is None else quantile
 
     def __eq__(self, __o: object) -> bool:
         return super().__eq__(__o) and self._quantile == __o._quantile
 
     def __hash__(self) -> int:
-        return hash((self.__class__.__name__, self.inplace, self._reduction_axes, tuple(self._quantile)))
+        return hash((self.__class__.__name__, self.inplace, self._axes, self._axes_mode, tuple(self._quantile)))
 
 
 class QuantileReducer(QuantileReducerBase):
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = x[0]
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         return fns.quantile(x, self._quantile, reduction_axes, keepdims=self._keepdims)
 
 
 class AbsQuantileReducer(QuantileReducerBase):
     def __init__(
         self,
-        reduction_axes: Optional[ReductionAxes] = None,
-        quantile: Optional[Union[float, list[float]]] = None,
+        axes: Optional[Axes] = None,
+        axes_mode: AxesMode = AxesMode.REDUCTION,
+        quantile: Optional[Union[float, tuple[float]]] = None,
         inplace: bool = False,
     ):
         quantile = (0.99,) if quantile is None else quantile
-        super().__init__(reduction_axes=reduction_axes, quantile=quantile, inplace=False)
+        super().__init__(axes, axes_mode, quantile)
 
     def _reduce_out_of_place(self, x: list[Tensor]) -> list[Tensor]:
         x = fns.abs(x[0])
-        reduction_axes = self._get_reduction_axes(x)
+        reduction_axes = determine_reduction_axes(x.ndim, self._axes, self._axes_mode)
         return fns.quantile(x, self._quantile, reduction_axes, keepdims=self._keepdims)
 
 
@@ -553,7 +604,7 @@ def __eq__(self, __o: object) -> bool:
         return super().__eq__(__o) and self._channel_axis == __o._channel_axis
 
     def __hash__(self) -> int:
-        return hash((self.__class__.__name__, self.inplace, self._reduction_axes, self._channel_axis))
+        return hash((self.__class__.__name__, self.inplace, self._axes, self._axes_mode, self._channel_axis))
 
 
 ##################################################

src/nncf/openvino/statistics/collectors.py

@@ -44,37 +44,37 @@
 
 class OVMinReducer(MinReducer):
     def get_inplace_fn(self):
-        return get_inplace_min_op(self._reduction_axes)
+        return get_inplace_min_op(self._axes)
 
 
 class OVMaxReducer(MaxReducer):
     def get_inplace_fn(self):
-        return get_inplace_max_op(self._reduction_axes, False)
+        return get_inplace_max_op(self._axes, False)
 
 
 class OVAbsMaxReducer(AbsMaxReducer):
     def get_inplace_fn(self):
-        return get_inplace_max_op(self._reduction_axes, True)
+        return get_inplace_max_op(self._axes, True)
 
 
 class OVMeanReducer(MeanReducer):
     def get_inplace_fn(self):
-        return get_inplace_mean_op(self._reduction_axes)
+        return get_inplace_mean_op(self._axes)
 
 
 class OVMeanVarianceReducer(MeanVarianceReducer):
     def get_inplace_fn(self):
-        return get_inplace_mean_var_op(self._reduction_axes)
+        return get_inplace_mean_var_op(self._axes)
 
 
 class OVMaxVarianceReducer(MaxVarianceReducer):
     def get_inplace_fn(self):
-        return get_inplace_max_var_op(self._reduction_axes)
+        return get_inplace_max_var_op(self._axes)
 
 
 class OVMeanAbsMaxReducer(MeanAbsMaxReducer):
     def get_inplace_fn(self):
-        return get_inplace_mean_max_op(self._reduction_axes, True)
+        return get_inplace_mean_max_op(self._axes, True)
 
 
 class OVShapeReducer(ShapeReducer):

src/nncf/quantization/algorithms/channel_alignment/openvino_backend.py

@@ -20,6 +20,7 @@
 from nncf.common.graph.layer_attributes import ConvolutionLayerAttributes
 from nncf.common.graph.transformations.commands import TargetType
 from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase
+from nncf.experimental.common.tensor_statistics.collectors import AxesMode
 from nncf.experimental.common.tensor_statistics.collectors import MedianAggregator
 from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
 from nncf.experimental.common.tensor_statistics.statistics import MinMaxTensorStatistic
@@ -81,7 +82,7 @@ def get_statistic_collector(
         reduction_axes, q: float, num_samples: int, inplace: bool
     ) -> TensorStatisticCollectorBase:
         tensor_collector = TensorCollector(MinMaxTensorStatistic)
-        quantile_reducer = OVQuantileReducer(reduction_axes, (q, 1 - q), inplace)
+        quantile_reducer = OVQuantileReducer(reduction_axes, AxesMode.REDUCTION, (q, 1 - q), inplace)
 
         for port_id, container_key in enumerate([MinMaxTensorStatistic.MIN_STAT, MinMaxTensorStatistic.MAX_STAT]):
             aggregator = MedianAggregator(num_samples=num_samples, aggregation_axes=(0, 1))

src/nncf/quantization/algorithms/min_max/algorithm.py

@@ -570,14 +570,12 @@ def _get_statistic_collector(
                 else:
                     quantile = 1 - params.quantile_outlier_prob
                 reducer = self._backend_entity.reducer_map[statistic_type](
-                    reduction_axes=reduction_axes, inplace=inplace, quantile=[quantile]
+                    axes=reduction_axes, inplace=inplace, quantile=[quantile]
                 )
             else:
                 if use_abs_max and statistic_type == StatisticsType.MAX:
                     statistic_type = StatisticsType.ABS_MAX
-                reducer = self._backend_entity.reducer_map[statistic_type](
-                    reduction_axes=reduction_axes, inplace=inplace
-                )
+                reducer = self._backend_entity.reducer_map[statistic_type](axes=reduction_axes, inplace=inplace)
 
             kwargs = {
                 "num_samples": num_samples,