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

model_compression_toolkit/core/common/network_editors/actions.py

@@ -20,9 +20,15 @@ from typing import Callable
 from mct_quantizers import QuantizationMethod
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.core.common.graph.base_node import BaseNode
 
 
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.common.graph.base_node import BaseNode
+from model_compression_toolkit.core.common.quantization.quantization_params_fn_selection import \
+    get_activation_quantization_params_fn, get_weights_quantization_params_fn
+from model_compression_toolkit.core.common.quantization.quantization_fn_selection import \
+    get_weights_quantization_fn
+
 _EditRule = namedtuple('EditRule', 'filter action')
 
 
@@ -58,13 +64,15 @@ class BaseAction(ABC):
     """
 
     @abstractmethod
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         """
         Apply an action on the node after matching the node with a node filter.
 
         Args:
             node: Node to apply the action on.
             graph: Graph to apply the action on.
+            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.)
 
         Returns:
             Node after action is applied.
@@ -87,13 +95,15 @@ class ChangeCandidatesWeightsQuantConfigAttr(BaseAction):
         self.kwargs = kwargs
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         """
         Change the attribute 'attr_name' in weights quantization config candidates with 'attr_value'.
 
         Args:
             node: Node object to change its quant_config.
             graph: Graph to apply the action on.
+            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.)
         Returns:
             The node after its weights' quantization config candidates have been modified.
         """
@@ -118,7 +128,7 @@ class ChangeFinalWeightsQuantConfigAttr(BaseAction):
         self.kwargs = kwargs
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         if node.final_weights_quantization_cfg is not None:
             for parameter_name, parameter_value in self.kwargs.items():
                 node.final_weights_quantization_cfg.set_quant_config_attr(parameter_name, parameter_value,
@@ -137,13 +147,17 @@ class ChangeCandidatesActivationQuantConfigAttr(BaseAction):
         """
         self.kwargs = kwargs
 
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         """
         Change the attribute 'attr_name' in activation quantization configuration candidates with 'attr_value'.
 
         Args:
             node: Node object to change its quant_config.
             graph: Graph to apply the action on.
+            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.)
+        Returns:q
+            The node after its activation quantization configuration candidates have been modified.
         """
         for nqc in node.candidates_quantization_cfg:
             for parameter_name, parameter_value in self.kwargs.items():
@@ -162,12 +176,55 @@ class ChangeFinalActivationQuantConfigAttr(BaseAction):
         """
         self.kwargs = kwargs
 
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         if node.final_activation_quantization_cfg is not None:
             for parameter_name, parameter_value in self.kwargs.items():
                 node.final_activation_quantization_cfg.set_quant_config_attr(parameter_name, parameter_value)
 
 
+class ChangeQuantizationParamFunction(BaseAction):
+    """
+    Class ChangeQuantizationParamFunction to change a node's weights/activations quantization params function.
+    """
+
+    def __init__(self,
+                 attr_name: str = None,
+                 activation_quantization_params_fn: Callable = None,
+                 weights_quantization_params_fn: Callable = None):
+        """
+        Init a ChangeQuantizationParamFunction object.
+
+        Args:
+            attr_name: The weights attribute's name to set the weights quantization params function for (if setting weights params).
+            activation_quantization_params_fn: a params function for a node's activations.
+            weights_quantization_params_fn: a params function for a node's weights.
+        """
+        self.activation_quantization_params_fn = activation_quantization_params_fn
+        self.weights_quantization_params_fn = weights_quantization_params_fn
+        self.attr_name = attr_name
+
+    def apply(self, node: BaseNode, graph, fw_info):
+        """
+        Change the node's weights/activations quantization params function.
+
+        Args:
+            node: Node object to change its quantization params function.
+            graph: Graph to apply the action on.
+            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.)
+
+        Returns:
+            The node after its quantization params function has been modified.
+        """
+        for nqc in node.candidates_quantization_cfg:
+            if self.activation_quantization_params_fn is not None:
+                nqc.activation_quantization_cfg.set_activation_quantization_params_fn(
+                    self.activation_quantization_params_fn)
+            if self.weights_quantization_params_fn is not None:
+                (nqc.weights_quantization_cfg.get_attr_config(self.attr_name)
+                 .set_weights_quantization_params_fn(self.weights_quantization_params_fn))
+
+
 class ChangeFinalActivationQuantizationMethod(BaseAction):
     """
     Class ChangeFinalActivationQuantizationMethod to change a node's weights/activations quantizer function.
@@ -183,19 +240,31 @@ class ChangeFinalActivationQuantizationMethod(BaseAction):
 
         self.activation_quantization_method = activation_quantization_method
 
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         """
         Change the node's activations quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
+            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.)
 
         Returns:
             The node after its quantization function has been modified.
         """
 
         if self.activation_quantization_method is not None and node.final_activation_quantization_cfg is not None:
+
+            activation_quantization_params_fn = get_activation_quantization_params_fn(
+                self.activation_quantization_method)
+
+            node.final_activation_quantization_cfg.set_activation_quantization_params_fn(
+                activation_quantization_params_fn)
+
+            activation_quantization_fn = fw_info.activation_quantizer_mapping.get(self.activation_quantization_method)
+
+            node.final_activation_quantization_cfg.set_activation_quantization_fn(activation_quantization_fn)
             node.final_activation_quantization_cfg.activation_quantization_method = self.activation_quantization_method
 
 
@@ -213,23 +282,38 @@ class ChangeCandidatesActivationQuantizationMethod(BaseAction):
         """
         self.activation_quantization_method = activation_quantization_method
 
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         """
         Change the node's activations quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
+            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.)
 
+        Returns:
+            The node after its quantization function has been modified.
         """
         if self.activation_quantization_method is not None:
             for qc in node.candidates_quantization_cfg:
+                activation_quantization_params_fn = get_activation_quantization_params_fn(
+                    self.activation_quantization_method)
+
+                qc.activation_quantization_cfg.set_activation_quantization_params_fn(activation_quantization_params_fn)
+                activation_quantization_fn = fw_info.activation_quantizer_mapping.get(
+                    self.activation_quantization_method)
+
+                if activation_quantization_fn is None:
+                    Logger.critical('Unknown activation quantization method specified.')  # pragma: no cover
+
+                qc.activation_quantization_cfg.set_activation_quantization_fn(activation_quantization_fn)
                 qc.activation_quantization_cfg.activation_quantization_method = self.activation_quantization_method
 
 
 class ChangeFinalWeightsQuantizationMethod(BaseAction):
     """
-    Class ChangeFinalWeightsQuantizationMethod to change a node's weights/activations quantizer method.
+    Class ChangeFinalWeightsQuantizationMethod to change a node's weights/activations quantizer function.
     """
 
     def __init__(self, attr_name: str, weights_quantization_method=None):
@@ -244,19 +328,36 @@ class ChangeFinalWeightsQuantizationMethod(BaseAction):
         self.weights_quantization_method = weights_quantization_method
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph):
+    def apply(self, node: BaseNode, graph, fw_info):
         """
         Change the node's weights quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
+            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.)
 
+        Returns:
+            The node after its quantization function has been modified.
         """
 
         if self.weights_quantization_method is not None and node.final_weights_quantization_cfg is not None:
-            attr_config = node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
-            attr_config.weights_quantization_method = self.weights_quantization_method
+
+            weights_quantization_params_fn = get_weights_quantization_params_fn(self.weights_quantization_method)
+
+            (node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
+             .set_weights_quantization_params_fn(weights_quantization_params_fn))
+
+            weights_quantization_fn = get_weights_quantization_fn(self.weights_quantization_method)
+
+            if weights_quantization_fn is None:
+                Logger.critical('Unknown weights quantization method specified.')  # pragma: no cover
+
+            (node.final_weights_quantization_cfg.get_attr_config(self.attr_name)
+             .set_weights_quantization_fn(weights_quantization_fn))
+            node.final_weights_quantization_cfg.get_attr_config(self.attr_name).weights_quantization_method = \
+                self.weights_quantization_method
 
 
 class ChangeCandidatesWeightsQuantizationMethod(BaseAction):
@@ -275,13 +376,15 @@ class ChangeCandidatesWeightsQuantizationMethod(BaseAction):
         self.weights_quantization_method = weights_quantization_method
         self.attr_name = attr_name
 
-    def apply(self, node: BaseNode, graph: Graph):
+    def apply(self, node: BaseNode, graph: Graph, fw_info: FrameworkInfo):
         """
         Change the node's weights quantization function.
 
         Args:
             node: Node object to change its threshold selection function.
             graph: Graph to apply the action on.
+            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.)
 
         Returns:
             The node after its quantization function has been modified.
@@ -289,7 +392,18 @@ class ChangeCandidatesWeightsQuantizationMethod(BaseAction):
 
         if self.weights_quantization_method is not None:
             for qc in node.candidates_quantization_cfg:
+
+                weights_quantization_params_fn = get_weights_quantization_params_fn(self.weights_quantization_method)
+
                 attr_qc = qc.weights_quantization_cfg.get_attr_config(self.attr_name)
+                attr_qc.set_weights_quantization_params_fn(weights_quantization_params_fn)
+
+                weights_quantization_fn = get_weights_quantization_fn(self.weights_quantization_method)
+
+                if weights_quantization_fn is None:
+                    Logger.critical('Unknown weights quantization method specified.')  # pragma: no cover
+
+                attr_qc.set_weights_quantization_fn(weights_quantization_fn)
                 attr_qc.weights_quantization_method = self.weights_quantization_method
 
 
@@ -308,13 +422,15 @@ class ReplaceLayer(BaseAction):
         self.layer_type = layer_type
         self.get_params_and_weights_fn = get_params_and_weights_fn
 
-    def apply(self, node: BaseNode, graph: Graph):
+    def apply(self, node: BaseNode, graph: Graph, fw_info: FrameworkInfo):
         """
         Replacing node's layer type and configurations
 
         Args:
             node: Node object to replace or modify
             graph: Graph to apply the action on.
+            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.)
 
         Returns:
             The node after its layer functionality has been modified.

model_compression_toolkit/core/common/network_editors/edit_network.py

@@ -14,17 +14,20 @@
 # ==============================================================================
 from typing import List
 
+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.network_editors import EditRule
 
 
 def edit_network_graph(graph: Graph,
+                       fw_info: FrameworkInfo,
                        network_editor: List[EditRule]):
     """
     Apply a list of edit rules on a graph.
 
     Args:
         graph: The graph to edit.
+        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.)
         network_editor: List of edit rules to apply to the graph.
 
@@ -35,5 +38,5 @@ def edit_network_graph(graph: Graph,
     for edit_rule in network_editor:
         filtered_nodes = graph.filter(edit_rule.filter)
         for node in filtered_nodes:
-            edit_rule.action.apply(node, graph)
+            edit_rule.action.apply(node, graph, fw_info)
     # return graph

model_compression_toolkit/core/common/pruning/channels_grouping.py

@@ -26,14 +26,18 @@ class ChannelGrouping:
     based on their importance scores and SIMD group sizes.
     """
 
-    def __init__(self, prunable_nodes: List[BaseNode]):
+    def __init__(self,
+                 prunable_nodes: List[BaseNode],
+                 fw_info: FrameworkInfo):
         """
         Initializes the ChannelGrouping with necessary information.
 
         Args:
             prunable_nodes: List of nodes that can be pruned.
+            fw_info: Framework-specific information and utilities.
         """
         self.prunable_nodes = prunable_nodes
+        self.fw_info = fw_info
         # Store for each node a list of numpy arrays. Each numpy array represents the
         # indices of the channels in an SIMD group.
         self._simd_groups_indices = {}

model_compression_toolkit/core/common/pruning/greedy_mask_calculator.py

@@ -38,6 +38,7 @@ class GreedyMaskCalculator:
     """
     def __init__(self,
                  prunable_nodes: List[BaseNode],
+                 fw_info: FrameworkInfo,
                  simd_groups_scores: Dict[BaseNode, np.ndarray],
                  target_resource_utilization: ResourceUtilization,
                  graph: Graph,
@@ -47,6 +48,7 @@ class GreedyMaskCalculator:
         """
         Args:
             prunable_nodes (List[BaseNode]): Nodes that are eligible for pruning.
+            fw_info (FrameworkInfo): Framework-specific information and utilities.
             simd_groups_scores (Dict[BaseNode, np.ndarray]): Importance scores for each SIMG group in a prunable node.
             target_resource_utilization (ResourceUtilization): The target resource utilization to achieve.
             graph (Graph): The computational graph of the model.
@@ -55,6 +57,7 @@ class GreedyMaskCalculator:
             simd_groups_indices (Dict[BaseNode, List[List[int]]]): Indices of SIMD groups in each node.
         """
         self.prunable_nodes = prunable_nodes
+        self.fw_info = fw_info
         self.target_resource_utilization = target_resource_utilization
         self.graph = graph
         self.fw_impl = fw_impl
@@ -64,11 +67,14 @@ class GreedyMaskCalculator:
         self.simd_groups_scores = simd_groups_scores
 
         self.oc_pruning_mask = PerSIMDGroupMask(prunable_nodes=prunable_nodes,
+                                                fw_info=fw_info,
                                                 simd_groups_indices=simd_groups_indices)
 
         self.memory_calculator = MemoryCalculator(graph=graph,
+                                                  fw_info=fw_info,
                                                   fw_impl=fw_impl)
 
+
     def get_mask(self) -> Dict[BaseNode, np.ndarray]:
         """
         Retrieves the current pruning mask for each prunable node.

model_compression_toolkit/core/common/pruning/importance_metrics/lfh_importance_metric.py

@@ -38,7 +38,8 @@ class LFHImportanceMetric(BaseImportanceMetric):
                  graph: Graph,
                  representative_data_gen: Callable,
                  fw_impl: PruningFrameworkImplementation,
-                 pruning_config: PruningConfig):
+                 pruning_config: PruningConfig,
+                 fw_info: FrameworkInfo):
         """
         Initialize the LFHImportanceMetric instance.
 
@@ -47,11 +48,13 @@ class LFHImportanceMetric(BaseImportanceMetric):
             representative_data_gen (Callable): Function to generate representative data.
             fw_impl (PruningFrameworkImplementation): Implementation of pruning for the framework.
             pruning_config (PruningConfig): Configuration for pruning.
+            fw_info (FrameworkInfo): Framework-specific information.
         """
         self.float_graph = graph
         self.representative_data_gen = representative_data_gen
         self.fw_impl = fw_impl
         self.pruning_config = pruning_config
+        self.fw_info = fw_info
 
         # Initialize internal dictionaries for storing intermediate computations.
         self._entry_node_to_hessian_score = {}
@@ -155,7 +158,8 @@ class LFHImportanceMetric(BaseImportanceMetric):
             Dict[BaseNode, List[np.ndarray]]: Dictionary of entry nodes mapped to their SIMD group indices.
         """
         # Initialize channel grouping utility.
-        channel_grouping = ChannelGrouping(prunable_nodes=list(entry_node_to_score.keys()))
+        channel_grouping = ChannelGrouping(prunable_nodes=list(entry_node_to_score.keys()),
+                                           fw_info=self.fw_info)
 
         channel_grouping.group_scores_by_simd_groups(entry_node_to_score)
         grouped_indices = channel_grouping.simd_groups_indices
@@ -245,14 +249,20 @@ class LFHImportanceMetric(BaseImportanceMetric):
         Returns:
             tuple: A tuple containing the kernel attribute, the number of output channels, and the axis of the output channels.
         """
+        kernel_attr = self.fw_info.get_kernel_op_attributes(entry_node.type)
+        # Ensure only one kernel attribute exists for the given node.
+        if len(kernel_attr) != 1:
+            Logger.critical(f"Expected a single attribute but found multiple attributes ({len(kernel_attr)}) for node {entry_node}.")
+        kernel_attr = kernel_attr[0]
+
         # Retrieve and validate the axis index for the output channels.
-        oc_axis = entry_node.channel_axis.output
+        oc_axis, _ = self.fw_info.kernel_channels_mapping.get(entry_node.type)
         if oc_axis is None or int(oc_axis) != oc_axis:
             Logger.critical(f"Invalid output channel axis type for node {entry_node}: expected integer but got {oc_axis}.")
 
         # Get the number of output channels based on the kernel attribute and axis.
-        num_oc = entry_node.get_weights_by_keys(entry_node.kernel_attr).shape[oc_axis]
-        return entry_node.kernel_attr, num_oc, oc_axis
+        num_oc = entry_node.get_weights_by_keys(kernel_attr[0]).shape[oc_axis]
+        return kernel_attr, num_oc, oc_axis
 
     def _concatenate_tensors_by_indices(self,
                                         channels: List[np.ndarray],

model_compression_toolkit/core/common/pruning/mask/per_channel_mask.py

@@ -35,8 +35,9 @@ class MaskIndicator(Enum):
     REMAINED = 1
 
 
+
 class PerChannelMask:
-    def __init__(self, prunable_nodes: List[BaseNode]):
+    def __init__(self, prunable_nodes: List[BaseNode], fw_info: FrameworkInfo):
         """
         Initializes the PerChannelMask with prunable nodes and framework information.
         This class is responsible for maintaining and updating the pruning masks for each
@@ -45,8 +46,10 @@ class PerChannelMask:
 
         Args:
             prunable_nodes: List of nodes in the model that are subject to pruning.
+            fw_info: Framework-specific information required for pruning operations.
         """
         self.prunable_nodes = prunable_nodes
+        self.fw_info = fw_info
         self._mask = None  # Initialize the mask dictionary
         self._init_masks()  # Call to initialize masks for each prunable node
 
@@ -103,7 +106,8 @@ class PerChannelMask:
         Returns:
             int: Number of output channels for the node.
         """
-        oc_axis = node.channel_axis.output
-        num_oc = node.get_weights_by_keys(node.kernel_attr).shape[oc_axis]
+        kernel_attr = self.fw_info.get_kernel_op_attributes(node.type)[0]
+        oc_axis = self.fw_info.kernel_channels_mapping.get(node.type)[0]
+        num_oc = node.get_weights_by_keys(kernel_attr).shape[oc_axis]
         return num_oc
 

model_compression_toolkit/core/common/pruning/mask/per_simd_group_mask.py

@@ -24,10 +24,10 @@ from model_compression_toolkit.core.common.pruning.memory_calculator import Memo
 from model_compression_toolkit.core.common.pruning.pruning_framework_implementation import PruningFrameworkImplementation
 from model_compression_toolkit.logger import Logger
 
-
 class PerSIMDGroupMask:
     def __init__(self,
                  prunable_nodes: List[BaseNode],
+                 fw_info: FrameworkInfo,
                  simd_groups_indices: Dict[BaseNode, List[List[int]]]):
         """
         Initializes a mask calculator for SIMD groups in prunable nodes.
@@ -35,11 +35,13 @@ class PerSIMDGroupMask:
 
         Args:
             prunable_nodes: List of nodes that can be pruned.
+            fw_info: Framework-specific information.
             simd_groups_indices: A dictionary mapping each node to its SIMD groups' indices.
         """
         # Initialize the per-channel mask
-        self.per_channel_mask = PerChannelMask(prunable_nodes=prunable_nodes)
+        self.per_channel_mask = PerChannelMask(prunable_nodes=prunable_nodes, fw_info=fw_info)
         self.prunable_nodes = prunable_nodes
+        self.fw_info = fw_info
         self.simd_groups_indices = simd_groups_indices
         self._mask_simd = None  # Initialize the SIMD group mask dictionary
         self._init_masks()  # Initialize masks for each prunable node

model_compression_toolkit/core/common/pruning/memory_calculator.py

@@ -34,16 +34,18 @@ class MemoryCalculator:
     which is crucial for deploying models on memory-constrained devices or optimizing for computational efficiency.
     """
 
-    def __init__(self, graph: Graph, fw_impl: PruningFrameworkImplementation):
+    def __init__(self, graph: Graph, fw_info: FrameworkInfo, fw_impl: PruningFrameworkImplementation):
         """
         Initializes the MemoryCalculator with necessary information about the model's graph,
         framework-specific details, and pruning implementation.
 
         Args:
             graph (Graph): Computational graph of the model.
+            fw_info (FrameworkInfo): Contains framework-specific information.
             fw_impl (PruningFrameworkImplementation): Implementation details for pruning.
         """
         self.graph = graph
+        self.fw_info = fw_info
         self.fw_impl = fw_impl
 
     def get_pruned_graph_memory(self,
@@ -202,13 +204,19 @@ class MemoryCalculator:
             if node == section.exit_node:
                 return masks.get(section.entry_node)
 
+        kernel_attr = self.fw_info.get_kernel_op_attributes(node.type)
+        # Ensure only one kernel attribute exists for the given node.
+        if len(kernel_attr) != 1:
+            Logger.critical(f"Expected a single attribute, but found {len(kernel_attr)} attributes for node '{node}'. Ensure the node configuration is correct.")
+        kernel_attr = kernel_attr[0]
+
         # Retrieve and validate the axis index for the output channels.
-        ic_axis = node.channel_axis.input
+        _, ic_axis = self.fw_info.kernel_channels_mapping.get(node.type)
         if ic_axis is None or int(ic_axis) != ic_axis:
             Logger.critical(f"Invalid input channel axis type for node '{node}': expected integer but got '{ic_axis}'.")
 
         # Get the number of output channels based on the kernel attribute and axis.
-        num_ic = node.get_weights_by_keys(node.kernel_attr).shape[ic_axis]
+        num_ic = node.get_weights_by_keys(kernel_attr).shape[ic_axis]
         mask = np.ones(num_ic, dtype=bool)
         return mask
 
@@ -281,7 +289,7 @@ class MemoryCalculator:
             int: The total number of parameters in the node after pruning.
         """
         total_params = 0
-        attributes_and_oc_axis = self.fw_impl.attrs_oi_channels_info_for_pruning(node)
+        attributes_and_oc_axis = self.fw_impl.attrs_oi_channels_info_for_pruning(node, self.fw_info)
 
         # Iterate over the node's weights and apply pruning based on the masks.
         for w_attr, w in node.weights.items():
@@ -303,7 +311,7 @@ class MemoryCalculator:
             num_oc = np.sum(output_mask)
         else:
             # Get the node channel axis from framework info
-            channel_axis = self.fw_impl.default_output_channel_axis if node.out_channel_axis is None else node.out_channel_axis
+            channel_axis = self.fw_info.out_channel_axis_mapping.get(node.type)
             if channel_axis is None:
                 Logger.critical(f"The channel axis is undefined. Please ensure the channel axis is explicitly defined for node {node.type} in the framework info.")
 

model_compression_toolkit/core/common/pruning/prune_graph.py

@@ -27,6 +27,7 @@ from model_compression_toolkit.logger import Logger
 
 def build_pruned_graph(graph: Graph,
                        masks: Dict[BaseNode, np.ndarray],
+                       fw_info: FrameworkInfo,
                        fw_impl: FrameworkImplementation) -> Graph:
     """
     Prunes the provided graph according to the given pruning output-channels masks.
@@ -34,6 +35,7 @@ def build_pruned_graph(graph: Graph,
     Args:
         graph: The original computational graph to be pruned.
         masks: A dictionary mapping each prunable node to its pruning mask.
+        fw_info: Framework-specific information object.
         fw_impl: Framework-specific implementation object.
 
     Returns:
@@ -64,7 +66,8 @@ def build_pruned_graph(graph: Graph,
             section_mask = PruningSectionMask(entry_node_oc_mask=mask,
                                               exit_node_ic_mask=mask)
             pruning_section.apply_inner_section_mask(section_mask,
-                                                     fw_impl)
+                                                     fw_impl,
+                                                     fw_info)
 
     return graph_to_prune
 

model_compression_toolkit/core/common/pruning/pruner.py

@@ -40,6 +40,7 @@ class Pruner:
     """
     def __init__(self,
                  float_graph: Graph,
+                 fw_info: FrameworkInfo,
                  fw_impl: PruningFrameworkImplementation,
                  target_resource_utilization: ResourceUtilization,
                  representative_data_gen: Callable,
@@ -48,6 +49,7 @@ class Pruner:
         """
         Args:
             float_graph (Graph): The floating-point representation of the model's computation graph.
+            fw_info (FrameworkInfo): Contains metadata and helper functions for the framework.
             fw_impl (PruningFrameworkImplementation): Implementation of specific framework methods required for pruning.
             target_resource_utilization (ResourceUtilization): The target resource utilization to be achieved after pruning.
             representative_data_gen (Callable): Generator function for representative dataset used in pruning analysis.
@@ -55,6 +57,7 @@ class Pruner:
             target_platform_capabilities (FrameworkQuantizationCapabilities): Object encapsulating the capabilities of the target hardware platform.
         """
         self.float_graph = float_graph
+        self.fw_info = fw_info
         self.fw_impl = fw_impl
         self.target_resource_utilization = target_resource_utilization
         self.representative_data_gen = representative_data_gen
@@ -81,6 +84,7 @@ class Pruner:
         # Apply Greedy strategy to compute masks based on importance scores.
         if self.pruning_config.channels_filtering_strategy == ChannelsFilteringStrategy.GREEDY:
             mask_calculator = GreedyMaskCalculator(entry_nodes,
+                                                   self.fw_info,
                                                    self.simd_scores,
                                                    self.target_resource_utilization,
                                                    self.float_graph,
@@ -95,6 +99,7 @@ class Pruner:
         Logger.info("Start pruning graph...")
         _pruned_graph = build_pruned_graph(self.float_graph,
                                            self.per_oc_mask,
+                                           self.fw_info,
                                            self.fw_impl)
         return _pruned_graph
 
@@ -111,7 +116,7 @@ class Pruner:
         # Retrieve and initialize the importance metric.
         im = get_importance_metric(self.pruning_config.importance_metric, graph=self.float_graph,
                                    representative_data_gen=self.representative_data_gen, fw_impl=self.fw_impl,
-                                   pruning_config=self.pruning_config)
+                                   pruning_config=self.pruning_config, fw_info=self.fw_info)
         entry_node_to_simd_score, simd_groups_indices = im.get_entry_node_to_simd_score(entry_nodes)
         return entry_node_to_simd_score, simd_groups_indices