Correcting comments on pull requests

Author: gouda-youichi

model_compression_toolkit/core/common/quantization/bit_width_config.py

@@ -15,7 +15,6 @@
 from dataclasses import dataclass, field
 from typing import List, Union, Dict
 
-from model_compression_toolkit.constants import WEIGHTS_ATTRIBUTE, ACTIVATION_ATTRIBUTE
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.matchers.node_matcher import BaseNodeMatcher
 from model_compression_toolkit.logger import Logger
@@ -25,25 +24,25 @@
 @dataclass
 class ManualBitWidthSelection:
     """
-   Class to encapsulate the manual bit width selection configuration for a specific filter.
+    Class to encapsulate the manual bit width selection configuration for a specific filter.
 
-   Attributes:
+    Attributes:
        filter (BaseNodeMatcher): The filter used to select nodes for bit width manipulation.
        bit_width (int): The bit width to be applied to the selected nodes.
-   """
+    """
     filter: BaseNodeMatcher
     bit_width: int
 
 @dataclass
 class ManualWeightsBitWidthSelection(ManualBitWidthSelection):
     """
-   Class to encapsulate the manual weights bit width selection configuration for a specific filter.
+    Class to encapsulate the manual weights bit width selection configuration for a specific filter.
 
-   Attributes:
+    Attributes:
        filter (BaseNodeMatcher): The filter used to select nodes for bit width manipulation.
        bit_width (int): The bit width to be applied to the selected nodes.
        attr (str): The filtered node's attributes to apply bit-width manipulation to.
-   """
+    """
     attr: WeightAttrT
 
 @dataclass
@@ -53,7 +52,7 @@ class BitWidthConfig:
 
     Attributes:
         manual_activation_bit_width_selection_list (List[ManualBitWidthSelection]): A list of ManualBitWidthSelection objects for activation defining manual bit-width configurations.
-        manual_activation_bit_width_selection_list (List[ManualWeightsBitWidthSelection]): A list of ManualWeightsBitWidthSelection for weights objects defining manual bit-width configurations.
+        manual_weights_bit_width_selection_list (List[ManualWeightsBitWidthSelection]): A list of ManualWeightsBitWidthSelection for weights objects defining manual bit-width configurations.
     """
     manual_activation_bit_width_selection_list: List[ManualBitWidthSelection] = field(default_factory=list)
     manual_weights_bit_width_selection_list: List[ManualWeightsBitWidthSelection] = field(default_factory=list)
@@ -69,31 +68,35 @@ def set_manual_activation_bit_width(self,
             bit_widths (Union[List[int], int]): The bit widths to be applied to the selected nodes.
             If a single value is given it will be applied to all the filters
         """
-        _, bit_widths, filters = self._expand_to_list_filter_and_bit_width(filters, bit_widths)
+        if filters is None:
+            Logger.critical(f"The filters cannot be None.")
+        _, bit_widths, filters = self._expand_to_list(filters, bit_widths)
         for bit_width, filter in zip (bit_widths, filters):
             self.manual_activation_bit_width_selection_list += [ManualBitWidthSelection(filter, bit_width)]
 
     def set_manual_weights_bit_width(self,
                                         filters: Union[List[BaseNodeMatcher], BaseNodeMatcher],
                                         bit_widths: Union[List[int], int],
-                                        attrs: Union[List[str], str]
+                                        attrs: Union[List[WeightAttrT], WeightAttrT]
                                      ):
         """
         Add a manual bit-width selection for weights to the configuration.
 
         Args:
             filters (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
             bit_widths (Union[List[int], int]): The bit widths for specified by attrs to be applied to the selected nodes.
-            attrs (Union[List[str], str]): The filtered node's attributes to apply bit-width manipulation to.
+            attrs (Union[List[WeightAttrT], WeightAttrT]): The filtered node's attributes to apply bit-width manipulation to.
             If a single value is given it will be applied to all the filters
         """
-        attrs, bit_widths, filters = self._expand_to_list_filter_and_bit_width(filters, bit_widths, attrs)
+        if filters is None:
+            Logger.critical(f"The filters cannot be None.")
+        attrs, bit_widths, filters = self._expand_to_list(filters, bit_widths, attrs)
         for attr, bit_width, filter in zip (attrs, bit_widths, filters):
             self.manual_weights_bit_width_selection_list += [ManualWeightsBitWidthSelection(filter, bit_width, attr)]
 
     def get_nodes_to_manipulate_activation_bit_widths(self, graph: Graph) -> Dict:
         """
-        Retrieve nodes from the graph that need their bit-widths changed according to the manual bit-width selections.
+        Retrieve nodes from the graph that need their bit-widths for activation changed according to the manual bit-width selections.
 
         Args:
             graph (Graph): The graph containing the nodes to be filtered and manipulated.
@@ -102,12 +105,11 @@ def get_nodes_to_manipulate_activation_bit_widths(self, graph: Graph) -> Dict:
             Dict: A dictionary mapping nodes to their new bit-widths.
         """
         activation_nodes_to_change_bit_width = self._construct_node_to_new_bit_mapping(graph, self.manual_activation_bit_width_selection_list)
-
         return activation_nodes_to_change_bit_width
 
     def get_nodes_to_manipulate_weights_bit_widths(self, graph: Graph) -> Dict:
         """
-        Retrieve nodes from the graph that need their bit-widths changed according to the manual bit-width selections.
+        Retrieve nodes from the graph that need their bit-widths for weights changed according to the manual bit-width selections.
 
         Args:
             graph (Graph): The graph containing the nodes to be filtered and manipulated.
@@ -116,15 +118,22 @@ def get_nodes_to_manipulate_weights_bit_widths(self, graph: Graph) -> Dict:
             Dict: A dictionary mapping nodes to their new bit-widths.
         """
         weights_nodes_to_change_bit_width = self._construct_node_to_new_bit_mapping(graph, self.manual_weights_bit_width_selection_list)
-
         return weights_nodes_to_change_bit_width
 
+    @staticmethod
+    def _expand_to_list_core(
+            filters: Union[List[BaseNodeMatcher], BaseNodeMatcher],
+            vals: Union[List[Union[WeightAttrT, int]], Union[WeightAttrT, int]]) -> list:
+        """
+        Extend the length of vals to match the length of filters.
 
+        Args:
+            filters (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
+            vals Union[List[Union[WeightAttrT, int], Union[WeightAttrT, int]]]): The bit widths or The filtered node's attributes.
 
-    def _expand_to_list_core(
-            self,
-            filters: Union[List[BaseNodeMatcher]],
-            vals: Union[List[any], any]):
+        Returns:
+            list: Extended vals to match the length of filters.
+        """
         vals = [vals] if not isinstance(vals, list) else vals
         if len(vals) > 1 and len(vals) != len(filters):
             Logger.critical(f"Configuration Error: The number of provided bit_width values {len(vals)} "
@@ -134,18 +143,39 @@ def _expand_to_list_core(
             vals = [vals[0] for f in filters]
         return vals
 
-    def _expand_to_list_filter_and_bit_width(
-            self,
+    @staticmethod
+    def _expand_to_list(
             filters: Union[List[BaseNodeMatcher]],
             bit_widths: Union[List[int], int],
-            attrs: Union[List[str], str] = None):
-        filters = [filters] if not isinstance(filters, list) else filters
-        bit_widths = self._expand_to_list_core(filters, bit_widths)
-        attrs = self._expand_to_list_core(filters, attrs)
+            attrs: Union[List[WeightAttrT], WeightAttrT] = None) -> [List]:
+        """
+        Extend the length of filters, bit-widths and The filtered node's attributes to match the length of filters.
+
+        Args:
+            filters (Union[List[BaseNodeMatcher], BaseNodeMatcher]): The filters used to select nodes for bit-width manipulation.
+            bit_widths (Union[List[int], int]): The bit widths for specified by attrs to be applied to the selected nodes.
+            attrs (Union[List[WeightAttrT], WeightAttrT]): The filtered node's attributes to apply bit-width manipulation to.
 
+        Returns:
+            [List]: A List of extended input arguments.
+        """
+        filters = [filters] if not isinstance(filters, list) else filters
+        bit_widths = BitWidthConfig._expand_to_list_core(filters, bit_widths)
+        if attrs is not None:
+            attrs = BitWidthConfig._expand_to_list_core(filters, attrs)
         return attrs, bit_widths, filters
 
-    def _construct_node_to_new_bit_mapping(self, graph, manual_bit_width_selection_list):
+    @staticmethod
+    def _construct_node_to_new_bit_mapping(graph, manual_bit_width_selection_list) -> Dict:
+        """
+        Retrieve nodes from the graph that need their bit-widths changed according to the manual bit-width selections.
+
+        Args:
+            graph (Graph): The graph containing the nodes to be filtered and manipulated.
+
+        Returns:
+            Dict: A dictionary retrieved nodes from the graph.
+        """
         unit_nodes_to_change_bit_width = {}
         for manual_bit_width_selection in manual_bit_width_selection_list:
             filtered_nodes = graph.filter(manual_bit_width_selection.filter)
@@ -154,15 +184,23 @@ def _construct_node_to_new_bit_mapping(self, graph, manual_bit_width_selection_l
                     f"Node Filtering Error: No nodes found in the graph for filter {manual_bit_width_selection.filter.__dict__} "
                     f"to change their bit width to {manual_bit_width_selection.bit_width}.")
             for n in filtered_nodes:
-                if n.get_node_weights_attributes() is False:
-                    Logger.critical(f'The requested attribute to change the bit width for {n} is not existing.')
-                # check if a manual configuration exists for this node
-                if n in unit_nodes_to_change_bit_width:
-                    Logger.info(
-                        f"Node {n} has an existing manual bit width configuration of {unit_nodes_to_change_bit_width.get(n)}. A new manual configuration request of {manual_bit_width_selection.bit_width} has been received, and the previous value is being overridden.")
-                if isinstance(manual_bit_width_selection_list,  ManualBitWidthSelection):
+                if type(manual_bit_width_selection) is ManualBitWidthSelection:
+                    # check if a manual configuration exists for this node
+                    if n in unit_nodes_to_change_bit_width:
+                        Logger.info(
+                            f"Node {n} has an existing manual bit width configuration of {unit_nodes_to_change_bit_width.get(n)}." 
+                            f"A new manual configuration request of {manual_bit_width_selection.bit_width} has been received, and the previous value is being overridden.")
                     unit_nodes_to_change_bit_width.update({n: manual_bit_width_selection.bit_width})
-                elif isinstance(manual_bit_width_selection_list, ManualWeightsBitWidthSelection):
+                elif type(manual_bit_width_selection) is ManualWeightsBitWidthSelection:
+                    if len(n.get_node_weights_attributes()) == 0:
+                        Logger.critical(f'The requested attribute to change the bit width for {n} is not existing.')
+
+                    if n in unit_nodes_to_change_bit_width:
+                        if unit_nodes_to_change_bit_width[n][1] == manual_bit_width_selection.attr:
+                            Logger.info(
+                                f"Node {n} has an existing manual bit width configuration of {unit_nodes_to_change_bit_width.get(n)}." 
+                                f"A new manual configuration request of {manual_bit_width_selection.bit_width} has been received, and the previous value is being overridden.")
                     unit_nodes_to_change_bit_width.update({n: [manual_bit_width_selection.bit_width, manual_bit_width_selection.attr]})
-
+                else:
+                    Logger.critical(f'The type of manual_bit_width_selection_list must be ManualBitWidthSelection or ManualWeightsBitWidthSelection.')
         return unit_nodes_to_change_bit_width