mct-nightly 2.4.0.20250617.613__py3-none-any.whl → 2.4.0.20250619.621__py3-none-any.whl

model_compression_toolkit/core/analyzer.py

@@ -32,8 +32,7 @@ def analyzer_model_quantization(representative_data_gen: Callable,
                                 tb_w: TensorboardWriter,
                                 float_graph: Graph,
                                 quantized_graph: Graph,
-                                fw_impl: FrameworkImplementation,
-                                fw_info: FrameworkInfo):
+                                fw_impl: FrameworkImplementation):
     """
     Plot the cosine similarity of different points on the graph between the float and quantized
     graphs. Add them to the passed TensorboardWriter object and close all tensorboard writer open
@@ -45,14 +44,12 @@ def analyzer_model_quantization(representative_data_gen: Callable,
         float_graph: Graph of float model.
         quantized_graph: Graph of quantized model.
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
-        fw_info: Information needed for quantization about the specific framework.
 
     """
     if tb_w is not None:
         visual = NNVisualizer(float_graph,
                               quantized_graph,
-                              fw_impl=fw_impl,
-                              fw_info=fw_info)
+                              fw_impl=fw_impl)
         if not visual.has_compare_points():
             Logger.error(f'No comparing points were found to plot analyze similarity.')
         else:

model_compression_toolkit/core/common/back2framework/base_model_builder.py

@@ -28,20 +28,17 @@ class BaseModelBuilder(ABC):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
-                 fw_info: FrameworkInfo = None,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes of graph 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.
         """
 
         self.graph = graph
         self.append2output = append2output
-        self.fw_info = fw_info
         self.return_float_outputs = return_float_outputs
 
     @abstractmethod

model_compression_toolkit/core/common/framework_implementation.py

@@ -125,18 +125,16 @@ class FrameworkImplementation(ABC):
                       graph: Graph,
                       mode: ModelBuilderMode,
                       append2output: List[Any],
-                      fw_info: FrameworkInfo,
                       return_float_outputs: bool = False) -> Tuple:
         """
         Build a framework model from a graph.
-        The mode determines how the model should be build. append2output is a list of Nodes
+        The mode determines how the model should be built. append2output is a list of Nodes
         to set as the model outputs.
 
         Args:
             graph: Graph to build the model from it.
             mode: Mode for how to build the model.
             append2output: List of Nodes to set as the model's outputs.
-            fw_info: FrameworkInfo object with information about the specific framework's model
             return_float_outputs (bool): whether to return outputs before or after quantization nodes (default)
 
         Returns:
@@ -170,15 +168,13 @@ class FrameworkImplementation(ABC):
     @abstractmethod
     def shift_negative_correction(self,
                                   graph: Graph,
-                                  core_config: CoreConfig,
-                                  fw_info: FrameworkInfo) -> Graph:
+                                  core_config: CoreConfig) -> Graph:
         """
         Apply shift negative correction (SNC) on a graph.
 
         Args:
             graph: Graph to apply SNC on.
             core_config: Quantization configuration.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             Graph after SNC.
@@ -189,15 +185,13 @@ class FrameworkImplementation(ABC):
     @abstractmethod
     def compute_activation_bias_correction(self,
                                            graph: Graph,
-                                           quant_config: QuantizationConfig,
-                                           fw_info: FrameworkInfo) -> Graph:
+                                           quant_config: QuantizationConfig) -> Graph:
         """
         Compute activation bias correction on a graph.
 
         Args:
             graph: Graph to apply activation bias correction on.
             quant_config: QuantizationConfig of how the model should be quantized.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             Graph after activation bias correction computing.
@@ -207,30 +201,28 @@ class FrameworkImplementation(ABC):
 
     @abstractmethod
     def get_substitutions_channel_equalization(self,
-                                               quant_config: QuantizationConfig,
-                                               fw_info: FrameworkInfo) -> List[common.BaseSubstitution]:
+                                               quant_config: QuantizationConfig) -> List[common.BaseSubstitution]:
         """
         Return a list of the framework substitutions used for channel equalization.
 
         Args:
             quant_config: QuantizationConfig to determine which substitutions to return.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             A list of the framework substitutions used after we collect statistics.
         """
         raise NotImplementedError(f'{self.__class__.__name__} has to implement the '
-                             f'framework\'s get_substitutions_channel_equalization method.')  # pragma: no cover
+                                  f'framework\'s get_substitutions_channel_equalization method.')  # pragma: no cover
 
     @abstractmethod
-    def get_substitutions_prepare_graph(self, fw_info: FrameworkInfo = None) -> List[common.BaseSubstitution]:
+    def get_substitutions_prepare_graph(self) -> List[common.BaseSubstitution]:
         """
 
         Returns: A list of the framework substitutions used to prepare the graph.
 
         """
         raise NotImplementedError(f'{self.__class__.__name__} has to implement the '
-                             f'framework\'s get_substitutions_prepare_graph method.')  # pragma: no cover
+                                  f'framework\'s get_substitutions_prepare_graph method.')  # pragma: no cover
 
     @abstractmethod
     def get_substitutions_pre_statistics_collection(self, quant_config: QuantizationConfig) -> \
@@ -328,14 +320,12 @@ class FrameworkImplementation(ABC):
                              f'method.')  # pragma: no cover
 
     def get_node_prior_info(self, node: BaseNode,
-                            fw_info: FrameworkInfo,
                             graph: Graph) -> NodePriorInfo:
         """
         Get a NodePriorInfo object for a node.
 
         Args:
             node: Node to get its prior info.
-            fw_info: Framework specific information needed to create the prior info of the node.
             graph: Graph to check the next node type.
 
         Returns:
@@ -343,7 +333,7 @@ class FrameworkImplementation(ABC):
         """
 
         raise NotImplementedError(f'{self.__class__.__name__} has to implement the '
-                             f'framework\'s get_node_prior_info method.')  # pragma: no cover
+                                  f'framework\'s get_node_prior_info method.')  # pragma: no cover
 
     def count_node_for_mixed_precision_interest_points(self, node: BaseNode) -> bool:
         """
@@ -394,20 +384,18 @@ class FrameworkImplementation(ABC):
 
     @abstractmethod
     def get_node_mac_operations(self,
-                                node: BaseNode,
-                                fw_info: FrameworkInfo) -> float:
+                                node: BaseNode) -> float:
         """
         Gets the MAC operation count for a given operation.
 
         Args:
             node: A graph node that wraps the operation for which the MAC count is computed.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns: The MAC count of the operation
         """
 
         raise NotImplementedError(f'{self.__class__.__name__} has to implement the '
-                             f'framework\'s get_node_mac_operations method.')  # pragma: no cover
+                                  f'framework\'s get_node_mac_operations method.')  # pragma: no cover
 
     @abstractmethod
     def apply_second_moment_correction(self,

model_compression_toolkit/core/common/framework_info.py

@@ -16,16 +16,16 @@
 
 from collections.abc import Callable
 from enum import Enum
-from typing import Dict, Any, List
+from typing import Dict, Any, Tuple, NamedTuple
+from abc import ABC, abstractmethod
 
 from mct_quantizers import QuantizationMethod
-from model_compression_toolkit.defaultdict import DefaultDict
 
 
 # Default value to use for ops without kernel.
 # This is a weird default, but it's used all over the place, so for now only extract it to const so that it can be
 # referenced by variable instead of hard-coded.
-DEFAULT_KERNEL_ATTRIBUTES = [None]
+DEFAULT_KERNEL_ATTRIBUTE = None
 
 
 class ChannelAxis(Enum):
@@ -42,89 +42,83 @@ class ChannelAxis(Enum):
     NCHW = 1
 
 
-class FrameworkInfo:
+class ChannelAxisMapping(NamedTuple):
+    output: int
+    input: int
 
-    def __init__(self,
-                 activation_quantizer_mapping: Dict[QuantizationMethod, Callable],
-                 kernel_channels_mapping: DefaultDict,
-                 activation_min_max_mapping: Dict[str, tuple],
-                 layer_min_max_mapping: Dict[Any, tuple],
-                 kernel_ops_attributes_mapping: DefaultDict,
-                 out_channel_axis_mapping: DefaultDict):
-        """
-        A class to wrap all information about a specific framework the library needs to quantize a model.
-        Specifically, FrameworkInfo holds lists of layers by how they should be quantized, and multiple mappings such as
-        layer to it kernel channels indices, and a layer to its min/max values, etc.
-        The layers lists are divided into three groups:
-        kernel_ops: Layers that have coefficients and need to get quantized (e.g., Conv2D, Dense, etc.)
-        activation_ops: Layers that their outputs should get quantized (e.g., Add, ReLU, etc.)
-        no_quantization_ops:Layers that should not get quantized (e.g., Reshape, Transpose, etc.)
-
-        Args:
-            activation_quantizer_mapping (Dict[QuantizationMethod, Callable]): A dictionary mapping from QuantizationMethod to a quantization function.
-            kernel_channels_mapping (DefaultDict): Dictionary from a layer to a tuple of its kernel in/out channels indices.
-            activation_min_max_mapping (Dict[str, tuple]): Dictionary from an activation function to its min/max output values.
-            layer_min_max_mapping (Dict[Any, tuple]): Dictionary from a layer to its min/max output values.
-            kernel_ops_attributes_mapping (DefaultDict): Dictionary from a framework operator to a list of its weights attirbutes to quantize.
-            out_channel_axis_mapping (DefaultDict): Dictionary of output channels of the model's layers (for computing statistics per-channel).
-
-        Examples:
-            When quantizing a Keras model, if we want to quantize the kernels of Conv2D layers only, we can
-            set, and we know it's kernel out/in channel indices are (3, 2) respectivly:
-
-            >>> import tensorflow as tf
-            >>> kernel_ops = [tf.keras.layers.Conv2D]
-            >>> kernel_channels_mapping = DefaultDict({tf.keras.layers.Conv2D: (3,2)})
 
-            Then, we can create a FrameworkInfo object:
+class FrameworkInfo(ABC):
+    """
+    A class to wrap all information about a specific framework the library needs to quantize a model.
+    Specifically, FrameworkInfo holds lists of layers by how they should be quantized, and multiple mappings such as
+    layer to it kernel channels indices, and a layer to its min/max values, etc.
+    The layers lists are divided into three groups:
+    kernel_ops: Layers that have coefficients and need to get quantized (e.g., Conv2D, Dense, etc.)
+    activation_ops: Layers that their outputs should get quantized (e.g., Add, ReLU, etc.)
+    no_quantization_ops:Layers that should not get quantized (e.g., Reshape, Transpose, etc.)
+
+    Fields:
+        activation_quantizer_mapping (Dict[QuantizationMethod, Callable]): A dictionary mapping from QuantizationMethod to a quantization function.
+        kernel_channels_mapping (Dict): Dictionary from a layer to a tuple of its kernel in/out channels indices.
+        kernel_ops_attribute_mapping (Dict): Dictionary from a framework operator to its weight attribute to quantize.
+        out_channel_axis_mapping (Dict): Dictionary of output channels of the model's layers (for computing statistics per-channel).
+        _layer_min_max_mapping (Dict[Any, tuple]): Dictionary from a layer to its min/max output values.
 
-            >>> FrameworkInfo(kernel_channels_mapping, {}, {})
+    """
 
-            If an activation layer (tf.keras.layers.Activation) should be quantized and we know it's min/max outputs range in advanced, we can add it to activation_min_max_mapping for saving the statistics collection time. For example:
+    activation_quantizer_mapping: Dict[QuantizationMethod, Callable]
+    kernel_channels_mapping: Dict[Any, ChannelAxisMapping]
+    kernel_ops_attribute_mapping: Dict[Any, str]
+    out_channel_axis_mapping: Dict[Any, int]
+    _layer_min_max_mapping: Dict[Any, tuple]
 
-            >>> activation_min_max_mapping = {'softmax': (0, 1)}
-            >>> FrameworkInfo(kernel_channels_mapping, activation_min_max_mapping, {})
+    _default_channel_mapping = ChannelAxisMapping(None, None)
 
-            If a layer's activations should be quantized and we know it's min/max outputs range in advanced, we can add it to layer_min_max_mapping for saving the statistics collection time. For example:
+    @classmethod
+    def get_kernel_op_attribute(cls, node_type: Any) -> str:
+        """
+        Get attribute of a layer's weight to quantize.
 
-            >>> layer_min_max_mapping = {tf.keras.layers.Softmax: (0, 1)}
-            >>> FrameworkInfo(kernel_channels_mapping, activation_min_max_mapping, layer_min_max_mapping)
+        Args:
+            node_type: Layer to get its attribute.
 
+        Returns:
+            Attribute the layer has and should be quantized.
         """
+        return cls.kernel_ops_attribute_mapping.get(node_type, DEFAULT_KERNEL_ATTRIBUTE)
 
-        self.activation_quantizer_mapping = activation_quantizer_mapping
-        self.kernel_channels_mapping = kernel_channels_mapping
-        self.activation_min_max_mapping = activation_min_max_mapping
-        self.layer_min_max_mapping = layer_min_max_mapping
-        self.kernel_ops_attributes_mapping = kernel_ops_attributes_mapping
-        self.out_channel_axis_mapping = out_channel_axis_mapping
-
-    def get_kernel_op_attributes(self, node_type: Any) -> List[str]:
+    @classmethod
+    def is_kernel_op(cls, node_type: Any) -> bool:
         """
-        Get a list of attributes of a layer's weights to quantize.
+        Check is the node is a kernel operation.
 
         Args:
             node_type: Layer to get its attributes.
 
         Returns:
-            A list of attributes the layer has and should be quantized.
+            True if node type is a kernel operation, else False.
         """
-        attr_list = self.kernel_ops_attributes_mapping.get(node_type)
-        return attr_list
+        return node_type in cls.kernel_ops_attribute_mapping
 
-    def is_kernel_op(self, node_type: Any) -> bool:
+    @classmethod
+    def get_layer_min_max(cls, layer: Any, fw_attrs: Dict) -> Tuple[float, float]:
         """
-        Check is the node is a kernel operation.
-
+        Return layer min/max mapping the FrameworkInfo holds.
         Args:
-            node_type: Layer to get its attributes.
+            layer: A layer to check if has a min/max known values.
+            fw_attrs: framework attributes from framework layer.
 
         Returns:
-            True if node type is a kernel operation, else False.
+            Layer's min/max known values.
         """
-        return node_type in self.kernel_ops_attributes_mapping.keys()
 
-    def layers_has_min_max(self, layer: Any) -> bool:
+        if cls.layers_has_min_max(layer):
+            return cls._layer_min_max_mapping[layer]
+        else:
+            return None, None
+
+    @classmethod
+    def layers_has_min_max(cls, layer: Any) -> bool:
         """
         Check if a layer is in a layer to min/max mapping the FrameworkInfo holds.
         Args:
@@ -134,17 +128,60 @@ class FrameworkInfo:
             Whether a layer has a min/max known values or not.
         """
 
-        return layer in self.layer_min_max_mapping
+        return layer in cls._layer_min_max_mapping
 
-    def activation_has_min_max(self, activation_name: str) -> bool:
+    @classmethod
+    @abstractmethod
+    def get_kernel_channels(cls, node_type: Any) -> ChannelAxisMapping:
         """
-        Check if an activation layer has a min/max mapping.
+        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.
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    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:
-            activation_name: String of the activation function to check for its min/max values.
+            node_type: A node type.
 
         Returns:
-            Whether an activation layer has a min/max known values or not.
+            Node's output channel axis.
+
         """
+        pass
+
+
+# Pointer to current FrameworkInfo class.
+_current_framework_info: type[FrameworkInfo] = None
+
+
+def get_fw_info():
+    """
+    A common function to get the current FrameworkInfo class. Raises an error if the pointer wasn't initialized.
+
+    Returns: FrameworkInfo class.
+    """
+    assert _current_framework_info is not None, "fw_info isn't initialized."
+    assert issubclass(_current_framework_info, FrameworkInfo), "fw_info isn't initialized to a FrameworkInfo class."
+    return _current_framework_info
+
+
+def set_fw_info(fw_info: type[FrameworkInfo]):
+    """
+    A common function to set the current FrameworkInfo class. Raises an error if fw_info doesn't inherit from FrameworkInfo.
+
+    Args:
+        fw_info: Framework specific object implementing the FrameworkInfo.
+    """
+    global _current_framework_info
+    assert _current_framework_info in [None, _current_framework_info], "FrameworkInfo already initialized."
+    assert issubclass(fw_info, FrameworkInfo), "fw_info must inherit from FrameworkInfo."
 
-        return activation_name in self.activation_min_max_mapping
+    _current_framework_info = fw_info

model_compression_toolkit/core/common/graph/base_graph.py

@@ -23,7 +23,6 @@ import numpy as np
 
 from networkx.algorithms.dag import topological_sort
 
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.fusion.fusing_info import FusingInfo
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX, EDGE_SOURCE_INDEX
 from model_compression_toolkit.core.common.graph.edge import Edge, convert_to_edge
@@ -74,7 +73,6 @@ class Graph(nx.MultiDiGraph, GraphSearches):
                  input_nodes: List[BaseNode],
                  output_nodes: List[OutTensor],
                  edge_list: List[Edge],
-                 fw_info: FrameworkInfo = None,
                  **attr):
         """
         Args:
@@ -82,7 +80,6 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             input_nodes: List of input nodes the model
             output_nodes: List of output nodes of the model to a list of their output indices.
             edge_list: List of edges the graph has between nodes.
-            fw_info: FrameworkInfo object (needed for computing the graph's weights memory).
             **attr: Attributes to add to graph as key=value pairs.
         """
 
@@ -103,7 +100,6 @@ class Graph(nx.MultiDiGraph, GraphSearches):
                           e.sink_node,
                           **e.get_attributes())
         self.user_info = UserInformation()
-        self.fw_info = fw_info
 
     @property
     def skip_validation_check(self) -> bool:
@@ -124,16 +120,6 @@ class Graph(nx.MultiDiGraph, GraphSearches):
     def fusing_info(self, fusing_info: FusingInfo):
         self._fusing_info = fusing_info
 
-    def set_fw_info(self,
-                    fw_info: FrameworkInfo):
-        """
-        Set the graph's framework info.
-        Args:
-            fw_info: FrameworkInfo object.
-        """
-
-        self.fw_info = fw_info
-
     def set_fqc(self,
                 fqc: FrameworkQuantizationCapabilities):
         """
@@ -563,7 +549,6 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return output_edges
 
     def get_configurable_sorted_nodes_names(self,
-                                            fw_info: FrameworkInfo,
                                             include_reused_nodes: bool = False) -> List[str]:
         """
         Get a list of nodes' names that can be configured (namely, has one or
@@ -571,56 +556,49 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         order of the graph.
 
         Args:
-            fw_info: FrameworkInfo object with information about the specific framework's model.
             include_reused_nodes: Whether or not to include reused nodes (False by default).
 
         Returns: List of nodes' names that can be configured (namely, has one or
         more weight qc candidate) sorted topology.
 
         """
-        sorted_names = [n.name for n in self.get_configurable_sorted_nodes(fw_info=fw_info,
-                                                                           include_reused_nodes=include_reused_nodes)]
+        sorted_names = [n.name for n in self.get_configurable_sorted_nodes(include_reused_nodes=include_reused_nodes)]
         return sorted_names
 
     def get_weights_configurable_nodes(self,
-                                       fw_info: FrameworkInfo,
                                        include_reused_nodes: bool = False) -> List[BaseNode]:
         """
         Get a list of nodes that their weights can be configured (namely, has one or
         more weight qc candidate and their weights should be quantized).
 
         Args:
-            fw_info: FrameworkInfo object with information about the specific framework's model.
             include_reused_nodes: Whether to include reused nodes (False by default).
 
         Returns:
             A list of nodes that their weights can be configured (namely, has one or more weight qc candidate).
         """
         # configurability is only relevant for kernel attribute quantization
-        potential_conf_nodes = [n for n in list(self) if fw_info.is_kernel_op(n.type)]
+        potential_conf_nodes = [n for n in list(self) if n.is_kernel_op]
 
         def is_configurable(n):
-            kernel_attrs = fw_info.get_kernel_op_attributes(n.type)
-            return any(n.is_configurable_weight(attr) for attr in kernel_attrs) and (not n.reuse or include_reused_nodes)
+            return n.is_configurable_weight(n.kernel_attr) and (not n.reuse or include_reused_nodes)
 
         return [n for n in potential_conf_nodes if is_configurable(n)]
 
     def get_sorted_weights_configurable_nodes(self,
-                                              fw_info: FrameworkInfo,
                                               include_reused_nodes: bool = False) -> List[BaseNode]:
         """
         Get a list of sorted nodes that their weights can be configured (namely, has one or
         more weight qc candidate and their weights should be quantized).
 
         Args:
-            fw_info: FrameworkInfo object with information about the specific framework's model.
             include_reused_nodes: Whether to include reused nodes (False by default).
 
         Returns:
             A list of nodes that their weights can be configured (namely, has one or more weight qc candidate)
             sorted topologically.
         """
-        return self._sort_nodes_in_list(self.get_weights_configurable_nodes(fw_info, include_reused_nodes))
+        return self._sort_nodes_in_list(self.get_weights_configurable_nodes(include_reused_nodes))
 
     def get_activation_configurable_nodes(self) -> List[BaseNode]:
         """
@@ -644,7 +622,6 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         return self._sort_nodes_in_list(self.get_activation_configurable_nodes())
 
     def get_configurable_sorted_nodes(self,
-                                      fw_info: FrameworkInfo,
                                       include_reused_nodes: bool = False) -> List[BaseNode]:
         """
         Get a list of nodes that can be configured (namely, has one or
@@ -652,14 +629,13 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         The nodes are sorted according to the topological order of the graph.
 
         Args:
-            fw_info: fw_info: FrameworkInfo object with information about the specific framework's model.
             include_reused_nodes: Whether or not to include reused nodes (False by default).
 
         Returns:
              A list of nodes that can be configured (namely, has one or more qc candidate) sorted topology.
 
         """
-        weights_configurable_nodes = self.get_weights_configurable_nodes(fw_info, include_reused_nodes)
+        weights_configurable_nodes = self.get_weights_configurable_nodes(include_reused_nodes)
         activation_configurable_nodes = self.get_activation_configurable_nodes()
 
         # combine and remove duplications
@@ -684,7 +660,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
                 sorted_configurable_nodes.append(n)
         return sorted_configurable_nodes
 
-    def get_min_candidates_config(self, fw_info: FrameworkInfo) -> Dict[BaseNode, int]:
+    def get_min_candidates_config(self) -> Dict[BaseNode, int]:
         """
         Builds a minimal configuration.
         Note: we assume that a minimal configuration exists, i.e., each configurable node has exactly one candidate
@@ -697,26 +673,23 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         Returns:
             A dict from layer to an index of its minimal candidate.
         """
-        conf_sorted_nodes = self.get_configurable_sorted_nodes(fw_info)
+        conf_sorted_nodes = self.get_configurable_sorted_nodes()
         return {n: n.find_min_candidate_index() for n in conf_sorted_nodes}
 
-    def get_max_candidates_config(self, fw_info: FrameworkInfo) -> Dict[BaseNode, int]:
+    def get_max_candidates_config(self) -> Dict[BaseNode, int]:
         """
         Builds a maximal configuration.
         Note: we assume that a maximal configuration exists, i.e., each configurable node has exactly one candidate
             with maximal n_bits (in both weight and activation if both are quantized, or in the relevant one if only
             one of them is quantized)
 
-        Args:
-            fw_info: fw_info: FrameworkInfo object with information about the specific framework's model.
-
         Returns:
             A dict from layer to an index of its maximal candidate.
         """
-        conf_sorted_nodes = self.get_configurable_sorted_nodes(fw_info)
+        conf_sorted_nodes = self.get_configurable_sorted_nodes()
         return {n: n.find_max_candidate_index() for n in conf_sorted_nodes}
 
-    def get_final_weights_config(self, fw_info: FrameworkInfo) -> List[Tuple[BaseNode, int]]:
+    def get_final_weights_config(self) -> List[Tuple[BaseNode, int]]:
         """
         Gets the final number of bits for quantization of each weights' configurable layer.
 
@@ -726,9 +699,9 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         Returns: A list of pairs of (node type, node's weights quantization bitwidth).
 
         """
-        sorted_conf_weights = self.get_sorted_weights_configurable_nodes(fw_info)
+        sorted_conf_weights = self.get_sorted_weights_configurable_nodes()
         # a configurable node by definition has a kernel op
-        return [(n, n.final_weights_quantization_cfg.get_attr_config(self.fw_info.get_kernel_op_attributes(n.type)[0]).weights_n_bits)
+        return [(n, n.final_weights_quantization_cfg.get_attr_config(n.kernel_attr).weights_n_bits)
                 for n in sorted_conf_weights]
 
     def get_final_activation_config(self) -> List[Tuple[BaseNode, int]]:
@@ -846,7 +819,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             next_node = self.out_edges(next_node)[0].sink_node
 
             # If next_node is an exit node and has only one incoming edge, the topology is prunable.
-            if fw_impl.is_node_exit_node(next_node, entry_node, self.fw_info) and len(self.in_edges(next_node)) == 1:
+            if fw_impl.is_node_exit_node(next_node, entry_node) and len(self.in_edges(next_node)) == 1:
                 return True
 
             # If the next node is not an intermediate node or has more than one incoming/outgoing edge,
@@ -876,7 +849,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
 
         intermediate_nodes, exit_node = self._find_intermediate_and_exit_nodes(entry_node, fw_impl)
 
-        if not fw_impl.is_node_exit_node(exit_node, entry_node, self.fw_info):
+        if not fw_impl.is_node_exit_node(exit_node, entry_node):
             Logger.critical(f"Node {exit_node} is not a valid exit node for the pruning section starting with {entry_node}.")   # pragma: no cover
 
         return PruningSection(entry_node=entry_node,
@@ -897,7 +870,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
         """
         intermediate_nodes = []
         next_node = self.out_edges(entry_node)[0].sink_node
-        while not fw_impl.is_node_exit_node(next_node, entry_node, self.fw_info):
+        while not fw_impl.is_node_exit_node(next_node, entry_node):
             intermediate_nodes.append(next_node)
             next_node = self.out_edges(next_node)[0].sink_node