mct-nightly 2.4.0.20250925.543__py3-none-any.whl → 2.4.2.20250926.532__py3-none-any.whl

model_compression_toolkit/core/common/quantization/node_quantization_config.py

@@ -12,15 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Any, List, Dict, TYPE_CHECKING
+
+
+from typing import Callable, Any, List, Tuple, Union, Dict, TYPE_CHECKING
 from enum import Enum, auto
+import numpy as np
 
-from model_compression_toolkit.core.common.framework_info import ChannelAxisMapping
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_weights_quantization_fn
 from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.core.common.quantization.quantization_params_fn_selection import \
+    get_activation_quantization_params_fn, get_weights_quantization_params_fn
 
-from model_compression_toolkit.target_platform_capabilities.constants import POSITIONAL_ATTR
+from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig, \
+    QuantizationErrorMethod
+from model_compression_toolkit.target_platform_capabilities.constants import POS_ATTR
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import \
-    AttributeQuantizationConfig, OpQuantizationConfig
+    AttributeQuantizationConfig, \
+    OpQuantizationConfig
 
 if TYPE_CHECKING:
     from model_compression_toolkit.core.common.graph.base_node import WeightAttrT
@@ -38,7 +46,6 @@ class ActivationQuantizationMode(Enum):
     FLN_QUANT = auto()
     PRESERVE_QUANT = auto()
     NO_QUANT = auto()
-    FLN_NO_QUANT = auto()
 
 
 class BaseNodeQuantizationConfig(object):
@@ -59,11 +66,12 @@ class BaseNodeQuantizationConfig(object):
             kwargs: A dictionary with additional key arguments.
 
         """
+
         if hasattr(self, config_parameter_name):
             setattr(self, config_parameter_name, config_parameter_value)
         else:
-            raise AttributeError(
-                f"Parameter {config_parameter_name} could not be found in the node quantization config.")
+            Logger.warning(f"Parameter {config_parameter_name} could not be found in the node quantization config and "
+                           f"was not updated!")
 
     def __repr__(self) -> str:
         """
@@ -77,14 +85,29 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
     """
     Attributes for configuring the quantization of the activations of a node.
     """
-    def __init__(self, op_cfg: OpQuantizationConfig):
+    def __init__(self,
+                 qc: QuantizationConfig,
+                 op_cfg: OpQuantizationConfig,
+                 activation_quantization_fn: Callable,
+                 activation_quantization_params_fn: Callable
+                 ):
         """
 
         Args:
+            qc: QuantizationConfig to create the node's config from.
             op_cfg: OpQuantizationConfig of the node with quantizers types to use when creating node quantization configuration.
+            activation_quantization_fn: Function to use when quantizing the node's activations.
+            activation_quantization_params_fn: Function to use when computing the threshold for quantizing a node's activations.
         """
+
+        self.activation_quantization_fn = activation_quantization_fn
+        self.activation_quantization_params_fn = activation_quantization_params_fn
+        self.activation_quantization_params = {}
         self.activation_quantization_method = op_cfg.activation_quantization_method
+        self.activation_error_method = qc.activation_error_method
         self.activation_n_bits = op_cfg.activation_n_bits
+        self.relu_bound_to_power_of_2 = qc.relu_bound_to_power_of_2
+        self.activation_bias_correction_term = None
         if op_cfg.enable_activation_quantization and op_cfg.quantization_preserving:
             raise ValueError("An OpQuantizationConfig can't have both enable_activation_quantization and quantization_preserving enabled.")
         if op_cfg.enable_activation_quantization:
@@ -94,13 +117,15 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
         else:
             self.quant_mode = ActivationQuantizationMode.NO_QUANT
         self.signedness = op_cfg.signedness
-
-        self.activation_quantization_params = {}
-        # TODO: computed by compute_activation_bias_correction. Probably shouldnt be here.
-        self.activation_bias_correction_term = None
-        # Z-threshold is a global param from QuantizationConfig, however it can be overridden per node by NetworkEditor.
-        # Since activation qparams are re-computed in several places, it's easier to keep it here and update it once.
-        self.z_threshold = None
+        self.activation_channel_equalization = qc.activation_channel_equalization
+        self.input_scaling = qc.input_scaling
+        self.min_threshold = qc.min_threshold
+        self.l_p_value = qc.l_p_value
+        self.shift_negative_activation_correction = qc.shift_negative_activation_correction
+        self.z_threshold = qc.z_threshold
+        self.shift_negative_ratio = qc.shift_negative_ratio
+        self.shift_negative_threshold_recalculation = qc.shift_negative_threshold_recalculation
+        self.concat_threshold_update = qc.concat_threshold_update
 
     @property
     def enable_activation_quantization(self):
@@ -113,6 +138,65 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
     def fln_quantization(self):
         return self.quant_mode == ActivationQuantizationMode.FLN_QUANT
 
+    def quantize_node_output(self,
+                             tensors: Any) -> Any:
+        """
+
+        Args:
+            tensors: framework tensor/s
+
+        Returns:
+            Framework tensor/s after applying fake quantization.
+
+        """
+        fake_quant = self.activation_quantization_fn(self.activation_n_bits,
+                                                     self.activation_quantization_params)
+
+        if fake_quant is None:
+            Logger.critical(
+                "Layer is intended to be quantized, but the fake_quant function is None.")  # pragma: no cover
+
+        return fake_quant(tensors)
+
+    @property
+    def activation_error_method(self) -> QuantizationErrorMethod:
+        """
+        activation_error_method getter.
+        """
+        return self._activation_error_method
+
+    @activation_error_method.setter
+    def activation_error_method(self, value: QuantizationErrorMethod):
+        """
+        activation_error_method setter.
+
+        Args:
+            value: New activation_error_method to set to the node activation configuration.
+
+        """
+        self._activation_error_method = value
+        self.activation_quantization_params_fn = get_activation_quantization_params_fn(activation_quantization_method=self.activation_quantization_method)
+
+    def set_activation_quantization_fn(self, activation_quantization_fn: Callable):
+        """
+        Sets activation quantization function for the node.
+
+        Args:
+            activation_quantization_fn: Function for quantazing the activations.
+
+        """
+        self.activation_quantization_fn = activation_quantization_fn
+
+    def set_activation_quantization_params_fn(self, activation_quantization_params_fn:Callable):
+        """
+        Sets activation params function for the node.
+
+        Args:
+            activation_quantization_params_fn: Function for calculating activation params.
+
+        """
+        self.activation_quantization_params_fn = activation_quantization_params_fn
+
     def set_activation_quantization_param(self,
                                           activation_params: dict):
         """
@@ -122,7 +206,7 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
             activation_params: Dictionary that contains weight quantization params.
 
         """
-        assert self.quant_mode == ActivationQuantizationMode.QUANT or self.quant_mode == ActivationQuantizationMode.FLN_QUANT
+        assert self.quant_mode == ActivationQuantizationMode.QUANT
         for param_name, param_value in activation_params.items():
             self.activation_quantization_params[param_name] = param_value
 
@@ -139,16 +223,36 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
         if not isinstance(other, NodeActivationQuantizationConfig):
             return False  # pragma: no cover
 
-        return self.activation_quantization_method == other.activation_quantization_method and \
+        return self.activation_quantization_fn == other.activation_quantization_fn and \
+               self.activation_quantization_params_fn == other.activation_quantization_params_fn and \
+               self.activation_error_method == other.activation_error_method and \
+               self.activation_quantization_method == other.activation_quantization_method and \
                self.activation_n_bits == other.activation_n_bits and \
                self.quant_mode == other.quant_mode and \
-               self.signedness == other.signedness
+               self.activation_channel_equalization == other.activation_channel_equalization and \
+               self.input_scaling == other.input_scaling and \
+               self.min_threshold == other.min_threshold and \
+               self.l_p_value == other.l_p_value and \
+               self.shift_negative_activation_correction == other.shift_negative_activation_correction and \
+               self.z_threshold == other.z_threshold and \
+               self.shift_negative_ratio == other.shift_negative_ratio and \
+               self.shift_negative_threshold_recalculation == other.shift_negative_threshold_recalculation
 
     def __hash__(self):
-        return hash((self.activation_quantization_method,
+        return hash((self.activation_quantization_fn,
+                     self.activation_quantization_params_fn,
+                     self.activation_error_method,
+                     self.activation_quantization_method,
                      self.activation_n_bits,
                      self.quant_mode,
-                     self.signedness))
+                     self.activation_channel_equalization,
+                     self.input_scaling,
+                     self.min_threshold,
+                     self.l_p_value,
+                     self.shift_negative_activation_correction,
+                     self.z_threshold,
+                     self.shift_negative_ratio,
+                     self.shift_negative_threshold_recalculation))
 
 
 class WeightsAttrQuantizationConfig:
@@ -156,21 +260,65 @@ class WeightsAttrQuantizationConfig:
     Configuration for quantizing a weights attribute of a node.
     """
     def __init__(self,
+                 qc: QuantizationConfig,
                  weights_attr_cfg: AttributeQuantizationConfig,
-                 weights_channels_axis: ChannelAxisMapping = None):
+                 weights_channels_axis: Tuple[int, int] = None):
         """
 
         Args:
+            qc: QuantizationConfig to create the node's config from.
             weights_attr_cfg: AttributeQuantizationConfig with parameters to use when creating the node's attribute quantization config.
             weights_channels_axis: Axis to quantize a node's attribute when quantizing per-channel (if not quantizing per-channel than expecting None).
         """
+        self.weights_quantization_fn = get_weights_quantization_fn(weights_attr_cfg.weights_quantization_method)
+        self.weights_quantization_params_fn = get_weights_quantization_params_fn(weights_attr_cfg.weights_quantization_method)
         self.weights_channels_axis = weights_channels_axis
+        self.weights_quantization_params = {}
         self.weights_quantization_method = weights_attr_cfg.weights_quantization_method
+        self.weights_error_method = qc.weights_error_method
         self.weights_n_bits = weights_attr_cfg.weights_n_bits
         self.weights_per_channel_threshold = weights_attr_cfg.weights_per_channel_threshold
         self.enable_weights_quantization = weights_attr_cfg.enable_weights_quantization
+        self.l_p_value = qc.l_p_value
 
-        self.weights_quantization_params = {}
+    @property
+    def weights_error_method(self) -> QuantizationErrorMethod:
+        """
+        weights_error_method getter.
+        """
+        return self._weights_error_method
+
+    @weights_error_method.setter
+    def weights_error_method(self, value: QuantizationErrorMethod):
+        """
+        weights_error_method setter.
+
+        Args:
+            value: New weights_error_method to set to the node weights configuration.
+
+        """
+        self._weights_error_method = value
+        self.weights_quantization_params_fn = get_weights_quantization_params_fn(weights_quantization_method=self.weights_quantization_method)
+
+    def set_weights_quantization_fn(self, weights_quantization_fn: Callable):
+        """
+        Sets weights quantization function for the node.
+
+        Args:
+            weights_quantization_fn: Function for quantazing the weights.
+
+        """
+        self.weights_quantization_fn = weights_quantization_fn
+
+    def set_weights_quantization_params_fn(self, weights_quantization_params_fn: Callable):
+        """
+        Sets weights params function for the node.
+
+        Args:
+            weights_quantization_params_fn: Function for calculating the weights params.
+
+        """
+        self.weights_quantization_params_fn = weights_quantization_params_fn
 
     def set_weights_quantization_param(self,
                                        weights_params: dict):
@@ -185,6 +333,31 @@ class WeightsAttrQuantizationConfig:
         for param_name, param_value in weights_params.items():
             self.weights_quantization_params[param_name] = param_value
 
+    def calculate_and_set_weights_params(self, tensor_data: np.ndarray, min_threshold: float):
+        """
+        Args:
+            tensor_data: Tensor content as Numpy array.
+            min_threshold: A minimal threshold to set as quantization parameter.
+
+        Returns:
+            Recalculated weights quantization params from the kernel and channel axis.
+
+        """
+        assert self.enable_weights_quantization
+        assert not (self.weights_per_channel_threshold and self.weights_channels_axis is None), \
+            "Trying to calculate threshold per channel, channel axis in None."
+        if self.weights_quantization_params_fn is not None:
+            self.set_weights_quantization_param(
+                self.weights_quantization_params_fn(tensor_data,
+                                                    p=self.l_p_value,
+                                                    n_bits=self.weights_n_bits,
+                                                    per_channel=self.weights_per_channel_threshold and self.weights_channels_axis is not None,
+                                                    channel_axis=self.weights_channels_axis[0],  # output channel axis
+                                                    min_threshold=min_threshold)[0]  # Take only first output, the q-params, as axis is already chosen.
+            )
+        else:
+            self.set_weights_quantization_param({})
+
     def __eq__(self, other: Any) -> bool:
         """
         Compares the object to another object to find if they are equal.
@@ -198,18 +371,26 @@ class WeightsAttrQuantizationConfig:
         if not isinstance(other, WeightsAttrQuantizationConfig):
             return False  # pragma: no cover
 
-        return self.weights_channels_axis == other.weights_channels_axis and \
+        return self.weights_quantization_fn == other.weights_quantization_fn and \
+               self.weights_quantization_params_fn == other.weights_quantization_params_fn and \
+               self.weights_channels_axis == other.weights_channels_axis and \
+               self.weights_error_method == other.weights_error_method and \
                self.weights_quantization_method == other.weights_quantization_method and \
                self.weights_n_bits == other.weights_n_bits and \
                self.weights_per_channel_threshold == other.weights_per_channel_threshold and \
-               self.enable_weights_quantization == other.enable_weights_quantization
+               self.enable_weights_quantization == other.enable_weights_quantization and \
+               self.l_p_value == other.l_p_value
 
     def __hash__(self):
-        return hash((self.weights_channels_axis,
+        return hash((self.weights_quantization_fn,
+                     self.weights_quantization_params_fn,
+                     self.weights_channels_axis,
+                     self.weights_error_method,
                      self.weights_quantization_method,
                      self.weights_n_bits,
                      self.weights_per_channel_threshold,
-                     self.enable_weights_quantization))
+                     self.enable_weights_quantization,
+                     self.l_p_value))
 
 
 class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
@@ -217,19 +398,23 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
     Holding a mapping between the node's weights attributes and their quantization configurations,
     in addition to quantization parameters that are global for all attributes of the represented node.
     """
-    def __init__(self,
+    def __init__(self, qc: QuantizationConfig,
                  op_cfg: OpQuantizationConfig,
-                 weights_channels_axis: ChannelAxisMapping,
+                 weights_channels_axis: Tuple[int, int],
                  node_attrs_list: List[str]):
         """
 
         Args:
+            qc: QuantizationConfig to create the node's config from.
             op_cfg: OpQuantizationConfig of the node with quantizers types to use when creating node quantization configuration.
             weights_channels_axis: Axis to quantize a node's weights attribute when quantizing per-channel.
             node_attrs_list: A list of the node's weights attributes names.
 
         """
+        self.min_threshold = qc.min_threshold
         self.simd_size = op_cfg.simd_size
+        self.weights_second_moment_correction = qc.weights_second_moment_correction
+        self.weights_bias_correction = qc.weights_bias_correction
 
         # Initialize a quantization configuration for each of the node's attributes
         self.attributes_config_mapping = {}
@@ -241,7 +426,7 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                 # POS_ATTR string. If none are found, it indicates that no specific quantization config is defined for
                 # positional weights, so the default config will be used instead.
                 attrs_included_in_name = {k: v for k, v in op_cfg.attr_weights_configs_mapping.items() if
-                                          POSITIONAL_ATTR in k}
+                                          POS_ATTR in k}
 
                 if len(attrs_included_in_name) > 1:  # pragma: no cover
                     raise ValueError(f"Found multiple attribute in FQC OpConfig that are contained "
@@ -257,7 +442,8 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                     attr_cfg = list(attrs_included_in_name.values())[0]
 
                 # Register this attribute under the positional attributes config mapping.
-                self.pos_attributes_config_mapping[attr] = WeightsAttrQuantizationConfig(weights_attr_cfg=attr_cfg,
+                self.pos_attributes_config_mapping[attr] = WeightsAttrQuantizationConfig(qc=qc,
+                                                                                         weights_attr_cfg=attr_cfg,
                                                                                          weights_channels_axis=
                                                                                          weights_channels_axis)
             else:
@@ -274,16 +460,9 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                 else:
                     attr_cfg = list(attrs_included_in_name.values())[0]
 
-                self.attributes_config_mapping[attr] = WeightsAttrQuantizationConfig(weights_attr_cfg=attr_cfg,
+                self.attributes_config_mapping[attr] = WeightsAttrQuantizationConfig(qc=qc,
+                                                                                     weights_attr_cfg=attr_cfg,
                                                                                      weights_channels_axis=weights_channels_axis)
-        # TODO this is set by batch norm reconstruction substitution when folded batch norms are added back, to mark
-        #  the nodes that the correction should be applied to (for some nodes it gets disabled) and BNs removed.
-        #  The actual correction is only computed when it's applied in ptq, so it seems that both substitutions could
-        #  be unified, and no info need to pass between.
-        self.weights_second_moment_correction = None
-        # TODO: computed corrected bias is injected to the node config. Probably shouldn't be here. Also it can be
-        #  computed on the final config, instead of all candidates and then there is no need to save it at all.
-        self.bias_corrected = None
 
     def get_attr_config(self, attr_name: 'WeightAttrT') -> WeightsAttrQuantizationConfig:
         """
@@ -420,8 +599,8 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                 if hasattr(attr_cfg, config_parameter_name):
                     setattr(attr_cfg, config_parameter_name, config_parameter_value)
                 else:
-                    raise AttributeError(f"Parameter {config_parameter_name} could not be found in the node quantization config of "
-                                         f"weights attribute {attr_name}.")
+                    Logger.warning(f"Parameter {config_parameter_name} could not be found in the node quantization config of "
+                                   f"weights attribute {attr_name} and was not updated!")
             else:  # pragma: no cover
                 Logger.critical(f"Weights attribute {attr_name} could not be found to set parameter {config_parameter_name}.")
 
@@ -438,7 +617,10 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
         if not isinstance(other, NodeWeightsQuantizationConfig):
             return False  # pragma: no cover
 
-        return self.simd_size == other.simd_size and \
+        return self.min_threshold == other.min_threshold and \
+            self.simd_size == other.simd_size and \
+            self.weights_second_moment_correction == other.weights_second_moment_correction and \
+            self.weights_bias_correction == other.weights_bias_correction and \
             self.attributes_config_mapping.keys() == other.attributes_config_mapping.keys() and \
             all([self.attributes_config_mapping[k] == other.attributes_config_mapping[k]
                  for k in self.attributes_config_mapping.keys()]) and \
@@ -447,6 +629,9 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                  for k in self.pos_attributes_config_mapping.keys()])
 
     def __hash__(self):
-        return hash((self.simd_size,
+        return hash((self.min_threshold,
+                     self.simd_size,
+                     self.weights_second_moment_correction,
+                     self.weights_bias_correction,
                      frozenset(self.attributes_config_mapping),
                      frozenset(self.pos_attributes_config_mapping)))

model_compression_toolkit/core/common/quantization/quantization_config.py

@@ -90,6 +90,7 @@ class QuantizationConfig:
     shift_negative_activation_correction: bool = True
     activation_channel_equalization: bool = False
     z_threshold: float = math.inf
+    min_threshold: float = MIN_THRESHOLD
     l_p_value: int = 2
     linear_collapsing: bool = True
     residual_collapsing: bool = True

model_compression_toolkit/core/common/quantization/quantization_fn_selection.py

@@ -14,35 +14,15 @@
 # ==============================================================================
 
 from collections.abc import Callable
+from functools import partial
 
 from mct_quantizers import QuantizationMethod
-
-from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeActivationQuantizationConfig
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.quantization.quantizers.lut_kmeans_quantizer import lut_kmeans_quantizer
 from model_compression_toolkit.core.common.quantization.quantizers.uniform_quantizers import power_of_two_quantizer, \
     symmetric_quantizer, uniform_quantizer
 
 
-def get_activation_quantization_fn(activation_quantization_cfg: NodeActivationQuantizationConfig,
-                                   get_activation_quantization_fn_factory: Callable) -> Callable:
-    """
-    Get activation quantizer based on activation quantization configuration.
-
-    Args:
-        activation_quantization_cfg: activation quantization configuration.
-        get_activation_quantization_fn_factory: activation quantization functions factory.
-
-    Returns:
-        Activation quantizer that accepts a tensor and returns a quantized tensor.
-    """
-    quantizer_factory = get_activation_quantization_fn_factory(
-        activation_quantization_cfg.activation_quantization_method)
-    quantizer = quantizer_factory(activation_quantization_cfg.activation_n_bits,
-                                  activation_quantization_cfg.activation_quantization_params)
-    return quantizer
-
-
 def get_weights_quantization_fn(weights_quantization_method: QuantizationMethod) -> Callable:
     """
     Generate a function for weight quantization.

model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py

@@ -0,0 +1,78 @@
+# Copyright 2021 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# 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 collections.abc import Callable
+from functools import partial
+
+from mct_quantizers import QuantizationMethod
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.lut_kmeans_params import \
+    lut_kmeans_tensor, lut_kmeans_histogram
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import \
+    symmetric_selection_tensor, symmetric_selection_histogram
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.uniform_selection import \
+    uniform_selection_histogram, uniform_selection_tensor
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.power_of_two_selection import \
+    power_of_two_selection_tensor, power_of_two_selection_histogram
+
+
+def get_activation_quantization_params_fn(activation_quantization_method: QuantizationMethod) -> Callable:
+    """
+    Generate a function for finding activation quantization parameters.
+
+    Args:
+        activation_quantization_method: Which quantization method to use for activations.
+    Returns:
+        A function to find the quantization parameters.
+
+    """
+    if activation_quantization_method == QuantizationMethod.POWER_OF_TWO:
+        params_fn = power_of_two_selection_histogram
+    elif activation_quantization_method == QuantizationMethod.SYMMETRIC:
+        params_fn = symmetric_selection_histogram
+    elif activation_quantization_method == QuantizationMethod.UNIFORM:
+        params_fn = uniform_selection_histogram
+    elif activation_quantization_method == QuantizationMethod.LUT_POT_QUANTIZER:
+        params_fn = lut_kmeans_histogram
+    else:
+        Logger.critical(
+            f"No parameter function found for the specified quantization method: {activation_quantization_method}")  # pragma: no cover
+    return params_fn
+
+
+def get_weights_quantization_params_fn(weights_quantization_method: QuantizationMethod) -> Callable:
+    """
+    Generate a function for finding weights quantization parameters.
+
+    Args:
+        weights_quantization_method: Which quantization method to use for weights.
+    Returns:
+        A function to find the quantization parameters.
+
+    """
+    if weights_quantization_method == QuantizationMethod.POWER_OF_TWO:
+        params_fn = power_of_two_selection_tensor
+    elif weights_quantization_method == QuantizationMethod.SYMMETRIC:
+        params_fn = symmetric_selection_tensor
+    elif weights_quantization_method == QuantizationMethod.UNIFORM:
+        params_fn = uniform_selection_tensor
+    elif weights_quantization_method == QuantizationMethod.LUT_POT_QUANTIZER:
+        params_fn = partial(lut_kmeans_tensor, is_symmetric=False)
+    elif weights_quantization_method == QuantizationMethod.LUT_SYM_QUANTIZER:
+        params_fn = partial(lut_kmeans_tensor, is_symmetric=True)
+    else:
+        Logger.critical(
+            f"No parameter function found for the specified quantization method: {weights_quantization_method}")  # pragma: no cover
+    return params_fn

model_compression_toolkit/core/common/quantization/quantization_params_generation/__init__.py

@@ -12,12 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.power_of_two_selection import (
-    power_of_two_no_clipping_selection_min_max, power_of_two_selection_histogram, power_of_two_selection_tensor)
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.lut_kmeans_params import (
-    lut_kmeans_tensor, lut_kmeans_histogram)
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import (
-    symmetric_no_clipping_selection_min_max, symmetric_selection_histogram, symmetric_selection_tensor)
-from model_compression_toolkit.core.common.quantization.quantization_params_generation.uniform_selection import (
-    uniform_no_clipping_selection_min_max, uniform_selection_histogram, uniform_selection_tensor)
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.power_of_two_selection import power_of_two_no_clipping_selection_min_max, \
+    power_of_two_selection_histogram, power_of_two_selection_tensor
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.lut_kmeans_params import lut_kmeans_tensor
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.symmetric_selection import symmetric_no_clipping_selection_min_max
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.uniform_selection import uniform_no_clipping_selection_min_max
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.outlier_filter import z_score_filter