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

model_compression_toolkit/core/keras/back2framework/keras_model_builder.py

@@ -35,6 +35,8 @@ from typing import Any, Dict, List, Tuple, Callable
 from tensorflow.python.util.object_identity import Reference as TFReference
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core import common
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.core.keras.back2framework.instance_builder import OperationHandler
@@ -55,6 +57,7 @@ class KerasModelBuilder(BaseModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False,
                  wrapper: Callable = None,
                  get_activation_quantizer_holder_fn: Callable=None):
@@ -63,6 +66,7 @@ class KerasModelBuilder(BaseModelBuilder):
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
+            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
             wrapper: A function wrapper keras Layers.
             get_activation_quantizer_holder_fn: Function to retrieve a quantization holder for a node.
@@ -71,6 +75,7 @@ class KerasModelBuilder(BaseModelBuilder):
 
         super().__init__(graph,
                          append2output,
+                         fw_info,
                          return_float_outputs)
 
         # Build an OperationHandler to handle conversions from graph nodes to Keras operators.

model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py

@@ -36,6 +36,7 @@ from model_compression_toolkit.core.keras.mixed_precision.configurable_weights_q
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 
 
 class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
@@ -46,12 +47,14 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
+            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
@@ -59,6 +62,7 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
 
         super().__init__(graph,
                          append2output,
+                         fw_info,
                          return_float_outputs,
                          wrapper=self.mixed_precision_wrapper,
                          get_activation_quantizer_holder_fn=self.mixed_precision_activation_holder)
@@ -83,12 +87,13 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
             ValueError: if kernel attribute is quantized but not configurable.
         """
 
-        if n.kernel_attr is None or not n.is_weights_quantization_enabled(n.kernel_attr):
+        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
+        if kernel_attr is None or not n.is_weights_quantization_enabled(kernel_attr):
             return layer
-        if not n.is_configurable_weight(n.kernel_attr):  # pragma: no cover
+        if not n.is_configurable_weight(kernel_attr):  # pragma: no cover
             raise ValueError(f'Weight wrapper is not expected to be created for non-configurable weight of node {n}.')
-        wq = ConfigurableWeightsQuantizer(**self._get_weights_configurable_quantizer_kwargs(n, n.kernel_attr))
-        return KerasQuantizationWrapper(layer, weights_quantizers={n.kernel_attr: wq})
+        wq = ConfigurableWeightsQuantizer(**self._get_weights_configurable_quantizer_kwargs(n, kernel_attr))
+        return KerasQuantizationWrapper(layer, weights_quantizers={kernel_attr: wq})
 
     def _get_weights_configurable_quantizer_kwargs(self, n: BaseNode, attr: str) -> Dict[str, Any]:
         """
@@ -142,12 +147,13 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
         # activation number of bits (in reversed order).
         # since only kernel attribute is quantized in weights mixed precision,
         # if the node doesn't have a kernel attribute, we only sort by activation_n_bits.
-        n.sort_node_candidates()
+        n.sort_node_candidates(self.fw_info)
 
         max_candidate_idx = n.find_max_candidate_index()
+        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
         activation_quantizers = [ConfigurableActivationQuantizer(**{'node_q_cfg': node_q_cfg_candidates,
                                                                     'max_candidate_idx': max_candidate_idx,
-                                                                    'kernel_attr': n.kernel_attr})] \
+                                                                    'kernel_attr': kernel_attr})] \
                                  * num_of_outputs
 
         # Holder by definition uses a single quantizer for the activation quantization
@@ -175,7 +181,7 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
 
         # creating a mapping between graph nodes and model's layers for mixed precision configurability
         conf_node2layers = {n.name: self._find_layers_in_model_by_node(n, model.layers)
-                            for n in self.graph.get_configurable_sorted_nodes()}
+                            for n in self.graph.get_configurable_sorted_nodes(self.fw_info)}
 
         return model, user_info, conf_node2layers
 
@@ -225,7 +231,8 @@ class MixedPrecisionKerasModelBuilder(KerasModelBuilder):
 
         """
         # Only layers with kernel op are considered weights configurable
-        weights_quant = False if n.kernel_attr is None else n.is_weights_quantization_enabled(n.kernel_attr)
+        kernel_attr = self.fw_info.get_kernel_op_attributes(n.type)[0]
+        weights_quant = False if kernel_attr is None else n.is_weights_quantization_enabled(kernel_attr)
         act_quant = n.is_activation_quantization_enabled()
 
         if weights_quant and not act_quant:

model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py

@@ -14,11 +14,11 @@
 # ==============================================================================
 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 model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from tensorflow.python.util.object_identity import Reference as TFReference
 
 
@@ -30,17 +30,20 @@ class QuantizedKerasModelBuilder(KerasModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
+                 fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
+            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
         super().__init__(graph,
                          append2output,
+                         fw_info,
                          return_float_outputs)
 
     def _quantize_node_activations(self,
@@ -57,6 +60,4 @@ class QuantizedKerasModelBuilder(KerasModelBuilder):
             Output of the node.
 
         """
-        activation_quantizer = get_activation_quantization_fn(node.final_activation_quantization_cfg,
-                                                              get_activation_quantization_fn_factory)
-        return activation_quantizer(input_tensors)
+        return node.final_activation_quantization_cfg.quantize_node_output(input_tensors)

model_compression_toolkit/core/keras/default_framework_info.py

@@ -13,142 +13,102 @@
 # limitations under the License.
 # ==============================================================================
 
-import tensorflow as tf
 
-from typing import Tuple, Any, Dict
-from functools import wraps
+import tensorflow as tf
 
+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"):
-    from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU, Activation
+    from keras.src.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU
 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 model_compression_toolkit.constants import SOFTMAX_THRESHOLD, ACTIVATION
+    from keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU  # pragma: no cover
+
+from model_compression_toolkit.defaultdict import DefaultDict
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo, DEFAULT_KERNEL_ATTRIBUTES
+from mct_quantizers import QuantizationMethod
+from model_compression_toolkit.constants import SOFTMAX_THRESHOLD
 from model_compression_toolkit.core.keras.constants import SOFTMAX, LINEAR, RELU, SWISH, SIGMOID, IDENTITY, TANH, SELU, \
     KERNEL, DEPTHWISE_KERNEL, GELU
-
-
-class KerasInfo(FrameworkInfo):
-    """
-    Extra field defined to handle Activation layer functions:
-
-        _activation_min_max_mapping (Dict[str, tuple]): Dictionary from an activation function to its min/max output values.
-
-    """
-
-    """
-    Map each layer to it's weight attribute that should get quantized.
-    If a layer that is not listed here is queried, None is returned.
-    """
-    kernel_ops_attribute_mapping = {Conv2D: KERNEL,
-                                    DepthwiseConv2D: DEPTHWISE_KERNEL,
-                                    Dense: KERNEL,
-                                    Conv2DTranspose: KERNEL}
-
-    """
-    Map a layer to its kernel's output and input channels indices.
-    Map's values are tuples of (output_channel_index, input_channel_index).
-    Default value is returned for layers that are not included.
-    """
-    kernel_channels_mapping = {Conv2D: ChannelAxisMapping(3, 2),
-                               DepthwiseConv2D: ChannelAxisMapping(2, 2),
-                               Dense: ChannelAxisMapping(1, 0),
-                               Conv2DTranspose: ChannelAxisMapping(2, 3)}
-
-    """
-    Map a layer to its output channel axis.
-    Where axis=-1 is the last axis
-    """
-    out_channel_axis_mapping = {Conv2D: -1,
-                                DepthwiseConv2D: -1,
-                                Dense: -1,
-                                Conv2DTranspose: -1}
-
-    """
-    Map from an activation function name to its min/max output values (if known).
-    The values are used for tensor min/max values initialization.
-    """
-    _activation_min_max_mapping = {SOFTMAX: (0, SOFTMAX_THRESHOLD),
-                                   SIGMOID: (0, 1),
-                                   LINEAR: (None, None),
-                                   IDENTITY: (None, None),
-                                   TANH: (-1, 1),
-                                   SWISH: (-0.279, None),
-                                   RELU: (0, None),
-                                   SELU: (-1.76, None),
-                                   GELU: (-0.17, None),
-                                   }
-
-    """
-    Map from an Keras module to its min/max output values (if known).
-    The values are used for tensor min/max values initialization.
-    """
-    _layer_min_max_mapping = {Softmax: (0, SOFTMAX_THRESHOLD),
-                              ELU: (-1, None),
-                              tf.nn.silu: (-0.279, None),
-                              tf.nn.swish: (-0.279, None),
-                              tf.nn.sigmoid: (0, 1),
-                              tf.nn.tanh: (-1, 1),
-                              tf.nn.relu: (0, None),
-                              tf.nn.relu6: (0, None),
-                              tf.nn.gelu: (-0.17, None),
-                              tf.nn.elu: (-1, None),
-                              tf.nn.selu: (-1.76, None),
-                              tf.nn.softplus: (0, None),
-                              tf.nn.softmax: (0, SOFTMAX_THRESHOLD),
-                              }
-
-    @classmethod
-    def get_layer_min_max(cls, layer: Any, fw_attrs: Dict) -> Tuple[float, float]:
-        """
-        Return layer min/max mapping the FrameworkInfo holds.
-        Args:
-            layer: A layer to check if has a min/max known values.
-            fw_attrs: framework attributes from framework layer.
-
-        Returns:
-            Layer's min/max known values.
-        """
-
-        if cls.layers_has_min_max(layer):
-            return cls._layer_min_max_mapping[layer]
-        elif isinstance(layer, Activation) and fw_attrs[ACTIVATION] in cls._activation_min_max_mapping:
-            return cls._activation_min_max_mapping[fw_attrs[ACTIVATION]]
-        else:
-            return None, None
-
-    @classmethod
-    def get_kernel_channels(cls, node_type: Any) -> ChannelAxisMapping:
-        """
-        Returns node's channels mapping from kernel_channels_mapping or framework specific default value.
-        Args:
-            node_type: A node type
-
-        Returns:
-            Node's channels mapping.
-
-        """
-        return cls.kernel_channels_mapping.get(node_type, cls._default_channel_mapping)
-
-    @classmethod
-    def get_out_channel_axis(cls, node_type: Any):
-        """
-        Returns node's output channel mapping from out_channel_axis_mapping or framework specific default value.
-        Args:
-            node_type: A node type.
-
-        Returns:
-            Node's output channel axis.
-
-        """
-        return cls.out_channel_axis_mapping.get(node_type)
-
-
-def set_keras_info(func):
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        set_fw_info(KerasInfo)
-        return func(*args, **kwargs)
-    return wrapper
+from model_compression_toolkit.core.keras.quantizer.fake_quant_builder import power_of_two_quantization, symmetric_quantization, uniform_quantization
+
+"""
+Map each layer to a list of its' weights attributes that should get quantized.
+If a layer that is not listed here is queried, [None] is returned.
+"""
+KERNEL_ATTRIBUTES = DefaultDict({Conv2D: [KERNEL],
+                                 DepthwiseConv2D: [DEPTHWISE_KERNEL],
+                                 Dense: [KERNEL],
+                                 Conv2DTranspose: [KERNEL]}, DEFAULT_KERNEL_ATTRIBUTES)
+
+
+"""
+Map a layer to its kernel's output and input channels indices.
+Map's values are tuples of (output_channel_index, input_channel_index).
+Default value is returned for layers that are not included.
+"""
+DEFAULT_CHANNEL_AXIS_DICT = DefaultDict({Conv2D: (3, 2),
+                                         DepthwiseConv2D: (2, 2),
+                                         Dense: (1, 0),
+                                         Conv2DTranspose: (2, 3)}, (None, None))
+
+
+"""
+Map a layer to its output channel axis. 
+Where axis=-1 is the last axis
+"""
+DEFAULT_OUT_CHANNEL_AXIS_DICT = DefaultDict({Conv2D: -1,
+                                             DepthwiseConv2D: -1,
+                                             Dense: -1,
+                                             Conv2DTranspose: -1},
+                                            -1)
+
+
+"""
+Map from an activation function to its min/max output values (if known).
+The values are used for tensor min/max values initialization.
+"""
+ACTIVATION2MINMAX = {SOFTMAX: (0, SOFTMAX_THRESHOLD),
+                     SIGMOID: (0, 1),
+                     LINEAR: (None, None),
+                     IDENTITY: (None, None),
+                     TANH: (-1, 1),
+                     SWISH: (-0.279, None),
+                     RELU: (0, None),
+                     SELU: (-1.76, None),
+                     GELU: (-0.17, None),
+                     }
+
+"""
+Map from an Keras layer to its min/max output values (if known).
+The values are used for tensor min/max values initialization.
+"""
+LAYER2MINMAX = {Softmax: (0, SOFTMAX_THRESHOLD),
+                ELU: (-1, None),
+                tf.nn.silu: (-0.279, None),
+                tf.nn.swish: (-0.279, None),
+                tf.nn.sigmoid: (0, 1),
+                tf.nn.tanh: (-1, 1),
+                tf.nn.relu: (0, None),
+                tf.nn.relu6: (0, None),
+                tf.nn.gelu: (-0.17, None),
+                tf.nn.elu: (-1, None),
+                tf.nn.selu: (-1.76, None),
+                tf.nn.softplus: (0, None),
+                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}
+
+
+DEFAULT_KERAS_INFO = FrameworkInfo(ACTIVATION_QUANTIZER_MAPPING,
+                                   DEFAULT_CHANNEL_AXIS_DICT,
+                                   ACTIVATION2MINMAX,
+                                   LAYER2MINMAX,
+                                   KERNEL_ATTRIBUTES,
+                                   DEFAULT_OUT_CHANNEL_AXIS_DICT)

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

@@ -21,6 +21,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
 from model_compression_toolkit.core.common.substitutions.batchnorm_folding import BatchNormalizationFolding, BatchNormalizationForwardFolding
 from model_compression_toolkit.core.keras.constants import KERNEL, LINEAR, ACTIVATION, DEPTHWISE_KERNEL, BIAS, GAMMA, BETA, \
     MOVING_MEAN, MOVING_VARIANCE, EPSILON, USE_BIAS, LAYER_NAME, GROUPS
+from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 
 
 def batchnorm_folding_node_matchers() -> [BaseNode, BaseNode]:
@@ -76,7 +77,9 @@ def update_kernel_for_bn_folding_fn(conv_node: BaseNode,
     else:
         kernel = kernel * weights_scale.reshape((1, 1, 1, -1))
 
-    return kernel, conv_node.kernel_attr
+    kernel_name = DEFAULT_KERAS_INFO.get_kernel_op_attributes(conv_node.type)[0]
+
+    return kernel, kernel_name
 
 
 def update_weights_for_bn_forward_folding_fn(conv_node: BaseNode,
@@ -105,7 +108,9 @@ def update_weights_for_bn_forward_folding_fn(conv_node: BaseNode,
         bias_update = (kernel * bias_factor.reshape((1, 1, -1, 1))).sum(2)
         kernel = kernel * weights_scale.reshape((1, 1, -1, 1))
 
-    return kernel, bias + bias_update.flatten(), conv_node.kernel_attr
+    kernel_name = DEFAULT_KERAS_INFO.get_kernel_op_attributes(conv_node.type)[0]
+
+    return kernel, bias + bias_update.flatten(), kernel_name
 
 
 def get_kernel_hw_fn(kernel: np.ndarray) -> [int, int]:

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

@@ -27,6 +27,7 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph, BaseNo
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.constants import REUSE, REUSE_GROUP
+from model_compression_toolkit.core.keras.reader.node_builder import REUSED_IDENTIFIER
 from model_compression_toolkit.core.keras.constants import KERNEL, BIAS, USE_BIAS, NUM_HEADS, KEY_DIM, VALUE_DIM, \
     QUERY_SHAPE, KEY_SHAPE, VALUE_SHAPE, OUTPUT_SHAPE, ATTENTION_AXES, ACTIVATION, GROUPS, LINEAR, FILTERS, PADDING, \
     FUNCTION, DIMS, TARGET_SHAPE, F_STRIDED_SLICE, F_STACK, Q_KERNEL, Q_BIAS, K_KERNEL, K_BIAS, V_KERNEL, V_BIAS, \

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

@@ -17,13 +17,14 @@
 from tensorflow.keras.layers import InputLayer, Dense, DepthwiseConv2D, Conv2D, Conv2DTranspose, ZeroPadding2D
 from typing import List
 
-from model_compression_toolkit.core import common, QuantizationConfig
+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, WalkMatcher
+from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, EdgeMatcher, 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.common.quantization.quantization_params_generation.qparams_weights_computation import \
-    compute_weights_qparams
+from model_compression_toolkit.core.keras.constants import KERNEL
 from model_compression_toolkit.logger import Logger
 
 input_node = NodeOperationMatcher(InputLayer)
@@ -47,8 +48,7 @@ class BaseInputScaling(common.BaseSubstitution):
     """
 
     def __init__(self,
-                 matcher_instance,
-                 quant_cfg: QuantizationConfig):
+                 matcher_instance):
         """
         Matches: InputLayer -> (optional nodes) -> (Dense,Conv2D,DepthwiseConv2D,Conv2DTranspose)
         note: the optional nodes are nodes that don't affect the scaling (such as ZeroPadding)
@@ -56,11 +56,10 @@ class BaseInputScaling(common.BaseSubstitution):
         Create a substitution using different params which may affect the way this substitution is made.
         The substitution is looking for edges in the graph which are input layers connected to linear layers.
         Args:
-            matcher_instance: matcher instance of type WalkMatcher.
-            quant_cfg: quantization config.
+            matcher_instance: matcher instance of type WalkMatcher
+
         """
         super().__init__(matcher_instance=matcher_instance)
-        self.quant_cfg = quant_cfg
 
     def substitute(self,
                    graph: Graph,
@@ -98,21 +97,17 @@ class BaseInputScaling(common.BaseSubstitution):
             scale_factor = threshold_float / threshold
             graph.user_info.set_input_scale(1 / scale_factor)
 
-            w1_fixed = linear_layer.get_weights_by_keys(linear_layer.kernel_attr) * scale_factor
-            linear_layer.set_weights_by_keys(linear_layer.kernel_attr, w1_fixed)
+            kernel_attr = graph.fw_info.get_kernel_op_attributes(linear_layer.type)[0]
+
+            w1_fixed = linear_layer.get_weights_by_keys(kernel_attr) * scale_factor
+            linear_layer.set_weights_by_keys(kernel_attr, w1_fixed)
 
             graph.scale_stats_collector(input_layer, 1 / scale_factor)
 
             # After scaling weights may have different thresholds so it needs to be recalculated
             for nqc in linear_layer.candidates_quantization_cfg:
-                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,
-                                                      weights_error_method=self.quant_cfg.weights_error_method,
-                                                      l_p_value=self.quant_cfg.l_p_value,
-                                                      output_channels_axis=attr_cfg.weights_channels_axis.output)
-                attr_cfg.set_weights_quantization_param(w_params)
+                nqc.weights_quantization_cfg.get_attr_config(kernel_attr).calculate_and_set_weights_params(w1_fixed,
+                                                                                                           nqc.weights_quantization_cfg.min_threshold)
 
         return graph
 
@@ -122,15 +117,12 @@ class InputScaling(BaseInputScaling):
     Substitution extends BaseInputScaling to the case of Input-->Linear
     """
 
-    def __init__(self, quant_cfg: QuantizationConfig):
+    def __init__(self):
         """
         Initialize a ScaleEqualization object.
-
-        Args:
-            quant_cfg: quantization config.
         """
 
-        super().__init__(matcher_instance=INPUT_MATCHER, quant_cfg=quant_cfg)
+        super().__init__(matcher_instance=INPUT_MATCHER)
 
 
 class InputScalingWithPad(BaseInputScaling):
@@ -138,12 +130,9 @@ class InputScalingWithPad(BaseInputScaling):
     Substitution extends BaseInputScaling to the case of Input-->ZeroPadding-->Linear
     """
 
-    def __init__(self, quant_cfg: QuantizationConfig):
+    def __init__(self):
         """
         Initialize a ScaleEqualization object.
-
-        Args:
-            quant_cfg: quantization config.
         """
 
-        super().__init__(matcher_instance=INPUT_MATCHER_WITH_PAD, quant_cfg=quant_cfg)
+        super().__init__(matcher_instance=INPUT_MATCHER_WITH_PAD)

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

@@ -63,15 +63,17 @@ class ScaleEqualization(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig):
+                 quant_config: QuantizationConfig,
+                 fw_info: FrameworkInfo):
         """
         Initialize a ScaleEqualization object.
         Args:
             quant_config: Quantization configuration.
+            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, matcher_instance=MATCHER,
+        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -81,15 +83,17 @@ class ScaleEqualizationWithPad(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig):
+                 quant_config: QuantizationConfig,
+                 fw_info: FrameworkInfo):
         """
         Initialize a ScaleEqualizationWithPad object.
         Args:
             quant_config: Quantization configuration.
+            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_WITH_PAD,
+        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER_WITH_PAD,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -100,15 +104,17 @@ class ScaleEqualizationMidActivation(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig):
+                 quant_config: QuantizationConfig,
+                 fw_info: FrameworkInfo):
         """
         Initialize a ScaleEqualizationMidActivation object.
         Args:
             quant_config: Quantization configuration.
+            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_MID,
+        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER_MID,
                          kernel_str=KERNEL, bias_str=BIAS)
 
 
@@ -118,13 +124,15 @@ class ScaleEqualizationMidActivationWithPad(BaseScaleEqualization):
     """
 
     def __init__(self,
-                 quant_config: QuantizationConfig):
+                 quant_config: QuantizationConfig,
+                 fw_info: FrameworkInfo):
         """
         Initialize a ScaleEqualizationMidActivationWithPad object.
         Args:
             quant_config: Quantization configuration.
+            fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
             groups of layers by how they should be quantized, etc.)
         """
 
-        super().__init__(quant_config=quant_config, matcher_instance=MATCHER_MID_WITH_PAD,
+        super().__init__(quant_config=quant_config, fw_info=fw_info, matcher_instance=MATCHER_MID_WITH_PAD,
                          kernel_str=KERNEL, bias_str=BIAS)