mct-nightly 2.4.0.20250630.629__py3-none-any.whl → 2.4.0.20250701.185106__py3-none-any.whl

model_compression_toolkit/core/common/statistics_correction/compute_activation_bias_correction_of_graph.py

@@ -18,7 +18,7 @@ from typing import Any, Callable
 from model_compression_toolkit.core import QuantizationConfig
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantization_fn
 
 
 def get_previous_node_with_activation_quantization(linear_node: BaseNode,
@@ -67,7 +67,8 @@ def compute_activation_bias_correction(graph: Graph,
                                        fw_impl: FrameworkImplementation,
                                        linear_node: BaseNode,
                                        prev_node: BaseNode,
-                                       kernel_size: str) -> Graph:
+                                       kernel_size: str,
+                                       get_activation_quantization_fn_factory: Callable) -> Graph:
     """
     Compute the activation bias correction term, and store it in the final activation
     quantization configuration.
@@ -79,6 +80,7 @@ def compute_activation_bias_correction(graph: Graph,
         linear_node: Node to compute the activation bias correction for.
         prev_node: Node to compute the activation error caused by his activation quantization.
         kernel_size: The framework specific attribute name of the convolution layer's kernel size.
+        get_activation_quantization_fn_factory: activation quantization functions factory.
 
     Returns:
         Graph with activation bias correction term for each node.
@@ -105,7 +107,9 @@ def compute_activation_bias_correction(graph: Graph,
     float_centers = calculate_bin_centers(float_bins)
 
     # Quantize the bin edges and calculate the centers of the quantized bins
-    quant_bins = prev_node_act_quant_cfg.quantize_node_output(fw_impl.to_tensor(float_bins))
+    activation_quantizer = get_activation_quantization_fn(prev_node_act_quant_cfg,
+                                                          get_activation_quantization_fn_factory)
+    quant_bins = activation_quantizer(fw_impl.to_tensor(float_bins))
     quant_bins = fw_impl.to_numpy(quant_bins)
     quant_centers = calculate_bin_centers(quant_bins)
 
@@ -149,7 +153,8 @@ def compute_activation_bias_correction_of_graph(graph: Graph,
                                                 quant_config: QuantizationConfig,
                                                 fw_impl: FrameworkImplementation,
                                                 activation_bias_correction_node_matchers: Callable,
-                                                kernel_size: str) -> Graph:
+                                                kernel_size: str,
+                                                get_activation_quantization_fn_factory: Callable) -> Graph:
     """
     Compute the activation bias correction term for the graph.
 
@@ -159,7 +164,7 @@ def compute_activation_bias_correction_of_graph(graph: Graph,
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         activation_bias_correction_node_matchers: Function to match the layers for activation bias correction.
         kernel_size: The framework specific attribute name of the convolution layer's kernel size.
-
+        get_activation_quantization_fn_factory: activation quantization functions factory.
 
     Returns:
         Graph with activation bias correction term for each relevant node.
@@ -175,5 +180,6 @@ def compute_activation_bias_correction_of_graph(graph: Graph,
                                                            fw_impl=fw_impl,
                                                            linear_node=n,
                                                            prev_node=prev_node,
-                                                           kernel_size=kernel_size)
+                                                           kernel_size=kernel_size,
+                                                           get_activation_quantization_fn_factory=get_activation_quantization_fn_factory)
     return graph

model_compression_toolkit/core/common/statistics_correction/compute_bias_correction_of_graph.py

@@ -43,7 +43,7 @@ def compute_bias_correction_of_graph(graph: Graph,
     for n in graph.nodes:
         # Bias correction is computed based on the quantized kernel, so we need to get the specific kernel attribute
         # name out of all the weights attributes of the node.
-        if n.is_kernel_op:
+        if n.kernel_attr:
             if n.is_weights_quantization_enabled(n.kernel_attr):
                 # Bias correction is not applied to layers with constant inputs.
                 if n.has_positional_weights:

model_compression_toolkit/core/common/substitutions/batchnorm_reconstruction.py

@@ -124,7 +124,7 @@ class BatchNormalizationReconstruction(common.BaseSubstitution):
 
         bn_node.prior_info = copy.deepcopy(source_node.prior_info)
 
-        bn_node.candidates_quantization_cfg = copy.deepcopy(source_node.candidates_quantization_cfg)
+        bn_node.quantization_cfg = copy.deepcopy(source_node.quantization_cfg)
 
         for qc in bn_node.candidates_quantization_cfg:
             qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.NO_QUANT
@@ -139,7 +139,6 @@ class BatchNormalizationReconstruction(common.BaseSubstitution):
                     #  reconstructed node BN attributes need to be quantized and how.
                     qc.weights_quantization_cfg.set_attr_config(attr,
                                                                 WeightsAttrQuantizationConfig(
-                                                                    QuantizationConfig(),
                                                                     AttributeQuantizationConfig(
                                                                         enable_weights_quantization=False)))
 

model_compression_toolkit/core/common/substitutions/shift_negative_activation.py

@@ -16,21 +16,20 @@ import copy
 import numpy as np
 from typing import List, Tuple, Any, Callable
 
-from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.core.common.quantization.node_quantization_config import WeightsAttrQuantizationConfig, \
     ActivationQuantizationMode
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.core.common import FrameworkInfo, Graph, BaseNode
+from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.constants import THRESHOLD, SIGNED, SHIFT_NEGATIVE_NON_LINEAR_NUM_BITS
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
-from model_compression_toolkit.core.common.quantization.set_node_quantization_config import create_node_activation_qc, \
-    set_quantization_configs_to_node
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_activations_computation \
-    import get_activations_qparams
+    import compute_activation_qparams
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \
     _mse_error_histogram
 from model_compression_toolkit.core.common.quantization.quantization_params_generation import z_score_filter
+from model_compression_toolkit.quantization_preparation.load_fqc import set_quantization_configs_to_node, \
+    fetch_qc_options_for_node
 from model_compression_toolkit.target_platform_capabilities import QuantizationMethod, AttributeQuantizationConfig
 
 """
@@ -67,8 +66,7 @@ def op2d_bias_correction(op2d_node: BaseNode,
         # Add an attribute quantization configuration to the newly added bias attribute, with disabled quantization
         for qc in op2d_node.candidates_quantization_cfg:
             qc.weights_quantization_cfg.set_attr_config(bias_flag_str,
-                                                        WeightsAttrQuantizationConfig(QuantizationConfig(),
-                                                                                      AttributeQuantizationConfig(
+                                                        WeightsAttrQuantizationConfig(AttributeQuantizationConfig(
                                                                                           enable_weights_quantization=False)))
 
     # Each node adds a different noise due to the shifting. It depends on the
@@ -253,6 +251,7 @@ def shift_negative_function(graph: Graph,
                             padding_str: str,
                             bias_str: str,
                             bias_flag_str: str,
+                            get_activation_quantization_fn_factory: Callable,
                             zero_padding_node: BaseNode = None,
                             bypass_nodes: List = None,
                             params_search_quantization_fn: Callable = None
@@ -278,6 +277,7 @@ def shift_negative_function(graph: Graph,
         padding_str: The framework specific attribute name of the padding.
         bias_str: The framework specific attribute name of the bias.
         bias_flag_str: The framework specific attribute name of the bias flag.
+        get_activation_quantization_fn_factory: activation quantization functions factory.
         zero_padding_node: ZeroPadding2D node that may be in the graph before the linear layer.
         params_search_quantization_fn: Function to quantize np tensor using a framework (tf/torch) quantization method. Needed for better param_search estimating the expected loss.
 
@@ -327,13 +327,15 @@ def shift_negative_function(graph: Graph,
                 'float32')  # Change to type float32 to support tensorflow dtypes
             for _shift_value in _q_points:
                 _hist_bins = hist_bins.astype(np.float32) + _shift_value
-                fw_quant_fn = non_linear_node_cfg_candidate.activation_quantization_fn(non_linear_node_cfg_candidate.activation_n_bits,qparams)
+                quantizer_factory = get_activation_quantization_fn_factory(
+                    non_linear_node_cfg_candidate.activation_quantization_method)
+                fw_quant_fn = quantizer_factory(non_linear_node_cfg_candidate.activation_n_bits, qparams)
                 """
                 In SNC, when better shifting values are tested for better choice,
                 the histogram (which is a numpy object) is quantized using the non-linear node activation
                 quantization function (to estimate the expected mse comparing to the original histogram).
                 The quantization function is a framework function, which makes it fail since it
-                expects a fw tensor. The commmon part of SNC receives an argument which is a callable 
+                expects a fw tensor. The common part of SNC receives an argument which is a callable 
                 that receives two argument and returns one: it gets the fw activation quantization function
                 and the bins to quantize. The function (of each fw) responsible for doing (if needed) a preprocessing and postprocessing
                 to the bins which is a numpy object.
@@ -395,9 +397,7 @@ def shift_negative_function(graph: Graph,
 
     set_quantization_configs_to_node(node=add_node,
                                      graph=graph,
-                                     quant_config=core_config.quantization_config,
-                                     fqc=graph.fqc,
-                                     mixed_precision_enable=core_config.is_mixed_precision_enabled)
+                                     fqc=graph.fqc)
 
     update_fused_op_with_add(graph=graph,
                              non_linear_node=non_linear_node,
@@ -421,9 +421,7 @@ def shift_negative_function(graph: Graph,
         # Set quantization configuration to node, even though we do not quantize it:
         set_quantization_configs_to_node(node=pad_node,
                                          graph=graph,
-                                         quant_config=core_config.quantization_config,
-                                         fqc=graph.fqc,
-                                         mixed_precision_enable=core_config.is_mixed_precision_enabled)
+                                         fqc=graph.fqc)
 
         for candidate_qc in pad_node.candidates_quantization_cfg:
             candidate_qc.activation_quantization_cfg.quant_mode = ActivationQuantizationMode.NO_QUANT
@@ -448,7 +446,7 @@ def shift_negative_function(graph: Graph,
                     bypass_candidate_qc.activation_quantization_cfg.activation_quantization_params[SIGNED] = False
                     graph.shift_stats_collector(bypass_node, np.array(shift_value))
 
-    add_node_qco = add_node.get_qco(graph.fqc).quantization_configurations
+    add_node_qco = fetch_qc_options_for_node(add_node, graph.fqc).quantization_configurations
     add_supported_bitwidths = [c.activation_n_bits for c in add_node_qco]
     if original_non_linear_activation_nbits not in add_supported_bitwidths:
         raise ValueError(
@@ -456,18 +454,16 @@ def shift_negative_function(graph: Graph,
             f"bitwidth is {original_non_linear_activation_nbits}. Consider adapting the TPC so 'Add' will support the "
             f"same bitwidth as {non_linear_node.type} or disable shift negative correction.")
 
-    for op_qc_idx, candidate_qc in enumerate(add_node.candidates_quantization_cfg):
-        for attr in add_node.get_node_weights_attributes():
-            # TODO: do we not quantize the weights of this 'add' on purpose?
-            candidate_qc.weights_quantization_cfg.get_attr_config(attr).enable_weights_quantization = False
+    set_quantization_configs_to_node(add_node, graph, graph.fqc)
+    # TODO: do we not quantize the weights of this 'add' on purpose?
+    add_node.quantization_cfg.disable_weights_quantization()
 
-        candidate_qc.activation_quantization_cfg = create_node_activation_qc(core_config.quantization_config,
-                                                                             add_node_qco[op_qc_idx])
+    def update(c):
+        c.activation_quantization_cfg.activation_n_bits = original_non_linear_activation_nbits
+        c.activation_quantization_cfg.set_activation_quantization_param({THRESHOLD: activation_threshold,
+                                                                         SIGNED: False})
 
-        candidate_qc.activation_quantization_cfg.set_activation_quantization_param({THRESHOLD: activation_threshold,
-                                                                                    SIGNED: False})
-
-        candidate_qc.activation_quantization_cfg.activation_n_bits = original_non_linear_activation_nbits
+    add_node.quantization_cfg.update_all(update, remove_duplicates=True)
 
     # Add the new padding node to a fused op with the op2d.
     if pad_node:
@@ -476,11 +472,11 @@ def shift_negative_function(graph: Graph,
                                op2d_node=op2d_node)
 
     if non_linear_node_cfg_candidate.shift_negative_threshold_recalculation:
-        activation_param = get_activations_qparams(activation_quant_cfg=non_linear_node_cfg_candidate,
-                                                   nodes_prior_info=non_linear_node.prior_info,
-                                                   out_stats_container=graph.get_out_stats_collector(non_linear_node))
+        activation_param = compute_activation_qparams(activation_quant_cfg=non_linear_node_cfg_candidate,
+                                                      node_prior_info=non_linear_node.prior_info,
+                                                      out_stats_container=graph.get_out_stats_collector(non_linear_node))
 
-        assert activation_param.get(SIGNED) == False
+        assert activation_param.get(SIGNED) is False
         for candidate_qc in non_linear_node.candidates_quantization_cfg:
             candidate_qc.activation_quantization_cfg.set_activation_quantization_param(activation_param)
 
@@ -573,6 +569,7 @@ def apply_shift_negative_correction(graph: Graph,
                                     padding_str: str,
                                     bias_str: str,
                                     bias_flag_str: str,
+                                    get_activation_quantization_fn_factory: Callable,
                                     params_search_quantization_fn: Callable=None) -> Graph:
     """
     Apply the substitution even if the linear node is not immediately after
@@ -594,6 +591,9 @@ def apply_shift_negative_correction(graph: Graph,
         padding_str: The framework specific attribute name of the padding.
         bias_str: The framework specific attribute name of the bias.
         bias_flag_str: The framework specific attribute name of the bias flag.
+        get_activation_quantization_fn_factory: activation quantization functions factory.
+        params_search_quantization_fn: Function to quantize np tensor using a framework (tf/torch) quantization method. Needed for better param_search estimating the expected loss.
+
     Returns:
         Graph after applying shift negative on selected activations.
     """
@@ -601,9 +601,8 @@ def apply_shift_negative_correction(graph: Graph,
     nodes = list(graph.nodes())
     for n in nodes:
         # Skip substitution if QuantizationMethod is uniform.
-        node_qco = n.get_qco(graph.fqc)
-        if any([op_qc.activation_quantization_method is QuantizationMethod.UNIFORM
-                for op_qc in node_qco.quantization_configurations]):
+        if any(aqc.activation_quantization_cfg.activation_quantization_method == QuantizationMethod.UNIFORM
+               for aqc in n.candidates_quantization_cfg):
             continue
 
         if snc_node_types.apply(n):
@@ -625,6 +624,7 @@ def apply_shift_negative_correction(graph: Graph,
                                                 padding_str,
                                                 bias_str,
                                                 bias_flag_str,
+                                                get_activation_quantization_fn_factory,
                                                 zero_padding_node=pad_node,
                                                 bypass_nodes=bypass_nodes,
                                                 params_search_quantization_fn=params_search_quantization_fn)

model_compression_toolkit/core/common/visualization/tensorboard_writer.py

@@ -207,7 +207,7 @@ class TensorboardWriter(object):
             attr = dict()
             if n.final_activation_quantization_cfg is not None:
                 attr.update(n.final_activation_quantization_cfg.__dict__)
-            elif n.candidates_quantization_cfg is not None:
+            elif n.quantization_cfg is not None:
                 attr.update(n.get_unified_activation_candidates_dict())
             return attr
 
@@ -229,7 +229,7 @@ class TensorboardWriter(object):
             attr = dict()
             if n.final_weights_quantization_cfg is not None:
                 attr.update(n.final_weights_quantization_cfg.__dict__)
-            elif n.candidates_quantization_cfg is not None:
+            elif n.quantization_cfg is not None:
                 attr.update(n.get_unified_weights_candidates_dict())
             return attr
 
@@ -530,8 +530,6 @@ def init_tensorboard_writer() -> TensorboardWriter:
     Create a TensorBoardWriter object initialized with the logger dir path if it was set,
     or None otherwise.
 
-    Args:
-
     Returns:
         A TensorBoardWriter object.
     """

model_compression_toolkit/core/graph_prep_runner.py

@@ -16,22 +16,22 @@
 
 from typing import Callable, Any
 
-from model_compression_toolkit.core.common import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
-from model_compression_toolkit.core.common.fusion.fusing_info import FusingInfoGenerator
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.filter_nodes_candidates import filter_nodes_candidates
-from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG
+from model_compression_toolkit.core.common.quantization.quantization_config import DEFAULTCONFIG, \
+    QuantizationErrorMethod
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
-from model_compression_toolkit.core.common.quantization.set_node_quantization_config import \
-    set_quantization_configuration_to_graph
+from model_compression_toolkit.core.common.quantization.set_node_quantization_config import set_manual_bitwidth_config
 from model_compression_toolkit.core.common.substitutions.apply_substitutions import substitute
 from model_compression_toolkit.core.common.substitutions.linear_collapsing_substitution import \
     linear_collapsing_substitute
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import TensorboardWriter
+from model_compression_toolkit.quantization_preparation.load_fqc import load_fqc_configuration
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.framework_quantization_capabilities import \
     FrameworkQuantizationCapabilities
+from model_compression_toolkit.logger import Logger
 
 
 def graph_preparation_runner(in_model: Any,
@@ -112,6 +112,12 @@ def get_finalized_graph(initial_graph: Graph,
 
     Returns: Graph object that represents the model, after applying all required modifications to it.
     """
+    if quant_config.weights_error_method == QuantizationErrorMethod.HMSE:
+        if not running_gptq:
+            raise ValueError(f"The HMSE error method for parameters selection is only supported when running GPTQ "
+                             f"optimization due to long execution time that is not suitable for basic PTQ.")
+        Logger.warning("Using the HMSE error method for weights quantization parameters search. "
+                       "Note: This method may significantly increase runtime during the parameter search process.")
 
     ######################################
     # Graph substitution (prepare graph)
@@ -141,21 +147,26 @@ def get_finalized_graph(initial_graph: Graph,
     if tb_w is not None:
         tb_w.add_graph(transformed_graph, 'pre_statistics_collection_substitutions')
 
-    ######################################
-    # Add quantization configurations
-    ######################################
-    transformed_graph = set_quantization_configuration_to_graph(graph=transformed_graph,
-                                                                quant_config=quant_config,
-                                                                bit_width_config=bit_width_config,
-                                                                mixed_precision_enable=mixed_precision_enable,
-                                                                running_gptq=running_gptq)
-
-    ######################################
-    # Layer fusing
-    ######################################
-    fusing_info = FusingInfoGenerator(fqc.get_fusing_patterns()).generate_fusing_info(transformed_graph)
-    transformed_graph.fusing_info = fusing_info
-    transformed_graph.override_fused_node_activation_quantization_candidates()
+    transformed_graph = load_fqc_configuration(transformed_graph, fqc)
+
+    # filter candidates per manual config
+    if bit_width_config:
+        set_manual_bitwidth_config(graph, bit_width_config)
+
+    # TODO irena: load_fqc_configuration only loads config from tpc. Previously quant_config was read as well.
+    #  As a first stage we keep the attributes in internal configs and fill them manually from quant_config
+    #  not to break all the code at once. Eventually we need to handle quant_config directly, without injecting into candidates.
+    #  TODO 2: Also we adjust candidates for single precision, which we shouldn't do here.
+    def update(qc):
+        qc.activation_quantization_cfg.set_qc(quant_config)
+        qc.weights_quantization_cfg.set_qc(quant_config)
+        for attr_cfg in qc.weights_quantization_cfg.get_all_weight_attrs_configs().values():
+            attr_cfg.weights_error_method = quant_config.weights_error_method
+            attr_cfg.l_p_value = quant_config.l_p_value
+    for n in transformed_graph.nodes:
+        if not mixed_precision_enable:
+            n.quantization_cfg.candidates_quantization_cfg = [n.quantization_cfg.base_quantization_cfg]
+        n.quantization_cfg.update_all(update)
 
     ######################################
     # Channel equalization

model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py

@@ -14,9 +14,10 @@
 # ==============================================================================
 from typing import List
 
-from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantization_fn
+from model_compression_toolkit.core.keras.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
 from tensorflow.python.util.object_identity import Reference as TFReference
 
@@ -56,4 +57,6 @@ class QuantizedKerasModelBuilder(KerasModelBuilder):
             Output of the node.
 
         """
-        return node.final_activation_quantization_cfg.quantize_node_output(input_tensors)
+        activation_quantizer = get_activation_quantization_fn(node.final_activation_quantization_cfg,
+                                                              get_activation_quantization_fn_factory)
+        return activation_quantizer(input_tensors)

model_compression_toolkit/core/keras/default_framework_info.py

@@ -18,7 +18,6 @@ import tensorflow as tf
 from typing import Tuple, Any, Dict
 from functools import wraps
 
-from model_compression_toolkit.core.keras.quantizer.lut_fake_quant import activation_lut_kmean_quantizer
 from packaging import version
 
 if version.parse(tf.__version__) >= version.parse("2.13"):
@@ -26,11 +25,9 @@ if version.parse(tf.__version__) >= version.parse("2.13"):
 else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU, Activation  # pragma: no cover
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo, set_fw_info, ChannelAxisMapping
-from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.constants import SOFTMAX_THRESHOLD, ACTIVATION
 from model_compression_toolkit.core.keras.constants import SOFTMAX, LINEAR, RELU, SWISH, SIGMOID, IDENTITY, TANH, SELU, \
     KERNEL, DEPTHWISE_KERNEL, GELU
-from model_compression_toolkit.core.keras.quantizer.fake_quant_builder import power_of_two_quantization, symmetric_quantization, uniform_quantization
 
 
 class KerasInfo(FrameworkInfo):
@@ -103,14 +100,6 @@ class KerasInfo(FrameworkInfo):
                               tf.nn.softmax: (0, SOFTMAX_THRESHOLD),
                               }
 
-    """
-    Mapping from a QuantizationMethod to an activation quantizer function.
-    """
-    activation_quantizer_mapping = {QuantizationMethod.POWER_OF_TWO: power_of_two_quantization,
-                                    QuantizationMethod.SYMMETRIC: symmetric_quantization,
-                                    QuantizationMethod.UNIFORM: uniform_quantization,
-                                    QuantizationMethod.LUT_POT_QUANTIZER: activation_lut_kmean_quantizer}
-
     @classmethod
     def get_layer_min_max(cls, layer: Any, fw_attrs: Dict) -> Tuple[float, float]:
         """

model_compression_toolkit/core/keras/graph_substitutions/substitutions/input_scaling.py

@@ -18,13 +18,12 @@ from tensorflow.keras.layers import InputLayer, Dense, DepthwiseConv2D, Conv2D,
 from typing import List
 
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.graph.base_graph import Graph
-from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, EdgeMatcher, WalkMatcher
+from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, WalkMatcher
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
-from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
 from model_compression_toolkit.constants import THRESHOLD
-from model_compression_toolkit.core.keras.constants import KERNEL
+from model_compression_toolkit.core.common.quantization.quantization_params_generation.qparams_weights_computation import \
+    compute_weights_qparams
 from model_compression_toolkit.logger import Logger
 
 input_node = NodeOperationMatcher(InputLayer)
@@ -104,8 +103,12 @@ class BaseInputScaling(common.BaseSubstitution):
 
             # After scaling weights may have different thresholds so it needs to be recalculated
             for nqc in linear_layer.candidates_quantization_cfg:
-                nqc.weights_quantization_cfg.get_attr_config(linear_layer.kernel_attr).calculate_and_set_weights_params(w1_fixed,
-                                                                                                           nqc.weights_quantization_cfg.min_threshold)
+                attr_cfg = nqc.weights_quantization_cfg.get_attr_config(linear_layer.kernel_attr)
+                assert attr_cfg.enable_weights_quantization
+                w_params, _ = compute_weights_qparams(w1_fixed, attr_quant_config=attr_cfg,
+                                                      output_channels_axis=attr_cfg.weights_channels_axis.output,
+                                                      min_threshold=nqc.weights_quantization_cfg.min_threshold)
+                attr_cfg.set_weights_quantization_param(w_params)
 
         return graph
 

model_compression_toolkit/core/keras/graph_substitutions/substitutions/shift_negative_activation.py

@@ -34,6 +34,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
     NodeFrameworkAttrMatcher
 from model_compression_toolkit.core.common.substitutions.shift_negative_activation import \
     apply_shift_negative_correction
+from model_compression_toolkit.core.keras.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.core.keras.constants import KERNEL_SIZE, STRIDES, ACTIVATION, SWISH, \
     SELU, GELU, FUNCTION, ADD, PAD
 from model_compression_toolkit.core.keras.constants import NEGATIVE_SLOPE, PADDING, PAD_SAME, PAD_VALID, BIAS, USE_BIAS
@@ -252,5 +253,6 @@ def keras_apply_shift_negative_correction(graph: Graph,
                                            is_padding_node_and_node_has_padding,
                                            PADDING,
                                            BIAS,
-                                           USE_BIAS
+                                           USE_BIAS,
+                                           get_activation_quantization_fn_factory
                                            )

model_compression_toolkit/core/keras/hessian/weights_hessian_scores_calculator_keras.py

@@ -94,7 +94,7 @@ class WeightsHessianScoresCalculatorKeras(HessianScoresCalculatorKeras):
                 for i, ipt_node in enumerate(self.hessian_request.target_nodes):  # Per Interest point weights tensor
 
                     # Check if the target node's layer type is supported.
-                    if not ipt_node.is_kernel_op:
+                    if not ipt_node.kernel_attr:
                         Logger.critical(f"Hessian information with respect to weights is not supported for "
                                         f"{ipt_node.type} layers.")  # pragma: no cover
 

model_compression_toolkit/core/keras/mixed_precision/configurable_activation_quantizer.py

@@ -23,6 +23,7 @@ from model_compression_toolkit.core.common.mixed_precision.configurable_quantize
     verify_candidates_descending_order, init_activation_quantizers
 from model_compression_toolkit.core.common.quantization.candidate_node_quantization_config import \
     CandidateNodeQuantizationConfig
+from model_compression_toolkit.core.keras.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.logger import Logger
 
 import tensorflow as tf
@@ -67,7 +68,7 @@ class ConfigurableActivationQuantizer(BaseKerasInferableQuantizer):
             if qc.activation_quantization_cfg.quant_mode != node_q_cfg[0].activation_quantization_cfg.quant_mode:
                 Logger.critical("Unsupported configuration: Mixing candidates with differing activation quantization states (enabled/disabled).")  # pragma: no cover
 
-        self.activation_quantizers = init_activation_quantizers(self.node_q_cfg)
+        self.activation_quantizers = init_activation_quantizers(self.node_q_cfg, get_activation_quantization_fn_factory)
         self.active_quantization_config_index = max_candidate_idx  # initialize with first config as default
 
     def set_active_activation_quantizer(self, index: Optional[int]):

model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py

@@ -155,7 +155,7 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         """
 
         attributes_with_axis = {}
-        if node.is_kernel_op:
+        if node.kernel_attr:
             attributes_with_axis[node.kernel_attr] = (node.channel_axis.output, node.channel_axis.input)
 
             # Bias is a vector at the length of the number of output channels.

model_compression_toolkit/core/keras/quantization/activation_quantization_fn_factory.py

@@ -0,0 +1,47 @@
+# Copyright 2025 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 mct_quantizers import QuantizationMethod
+
+
+from model_compression_toolkit.core.keras.quantization.fake_quant_builder import power_of_two_quantization, \
+    symmetric_quantization, uniform_quantization
+from model_compression_toolkit.core.keras.quantization.lut_fake_quant import activation_lut_kmean_quantizer
+
+
+"""
+Mapping from a QuantizationMethod to an activation quantizer function.
+"""
+_activation_quantizer_factory_mapping = {
+    QuantizationMethod.POWER_OF_TWO: power_of_two_quantization,
+    QuantizationMethod.SYMMETRIC: symmetric_quantization,
+    QuantizationMethod.UNIFORM: uniform_quantization,
+    QuantizationMethod.LUT_POT_QUANTIZER: activation_lut_kmean_quantizer
+}
+
+
+def get_activation_quantization_fn_factory(quantization_method: QuantizationMethod) -> Callable[[int, dict], Callable]:
+    """
+    Get factory for activation quantizer.
+
+    Args:
+        quantization_method: quantization method for activation.
+
+    Returns:
+        Factory that accepts activation bitwidth and a dict of quantization params, and returns the quantizer.
+    """
+    return _activation_quantizer_factory_mapping[quantization_method]

model_compression_toolkit/core/keras/statistics_correction/keras_compute_activation_bias_correction_of_graph.py

@@ -25,7 +25,7 @@ else:
 
 from model_compression_toolkit.core import QuantizationConfig
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.keras.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core.common.statistics_correction.compute_activation_bias_correction_of_graph import \
@@ -60,5 +60,6 @@ def keras_compute_activation_bias_correction_of_graph(graph: Graph,
                                                         fw_impl=fw_impl,
                                                         activation_bias_correction_node_matchers=
                                                         activation_bias_correction_node_matchers,
-                                                        kernel_size=KERNEL_SIZE)
+                                                        kernel_size=KERNEL_SIZE,
+                                                        get_activation_quantization_fn_factory=get_activation_quantization_fn_factory)
     return graph

model_compression_toolkit/core/pytorch/back2framework/quantized_model_builder.py

@@ -17,9 +17,10 @@ from typing import List, Tuple
 
 import torch
 
-from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import get_activation_quantization_fn
+from model_compression_toolkit.core.pytorch.quantization.activation_quantization_fn_factory import get_activation_quantization_fn_factory
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder, \
     PytorchModel
@@ -60,7 +61,9 @@ class QuantizedPyTorchModel(PytorchModel):
         if node.is_activation_quantization_enabled():
             if isinstance(input_tensors, list):
                 input_tensors = torch.cat(input_tensors, dim=0)
-            return node.final_activation_quantization_cfg.quantize_node_output(input_tensors)
+            activation_quantizer = get_activation_quantization_fn(node.final_activation_quantization_cfg,
+                                                                  get_activation_quantization_fn_factory)
+            return activation_quantizer(input_tensors)
         return input_tensors