mct-nightly 2.1.0.20240731.414__tar.gz → 2.1.0.20240801.515__tar.gz

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.1.0.20240731.414
+Version: 2.1.0.20240801.515
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.1.0.20240731.414
+Version: 2.1.0.20240801.515
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.1.0.20240731.000414"
+__version__ = "2.1.0.20240801.000515"

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/framework_implementation.py

@@ -348,24 +348,20 @@ class FrameworkImplementation(ABC):
         raise NotImplemented(f'{self.__class__.__name__} have to implement the '
                              f'framework\'s count_node_for_mixed_precision_interest_points method.')  # pragma: no cover
 
-    def get_mp_node_distance_fn(self, layer_class: type,
-                             framework_attrs: Dict[str, Any],
-                             compute_distance_fn: Callable = None,
-                             axis: int = None,
-                             norm_mse: bool = False) -> Callable:
+    def get_mp_node_distance_fn(self, n: BaseNode,
+                                compute_distance_fn: Callable = None,
+                                norm_mse: bool = False) -> Tuple[Callable, int]:
         """
         A mapping between layers' types and a distance function for computing the distance between
         two tensors in mixed precision (for loss computation purposes). Returns a specific function if node of specific types is
         given, or a default (normalized MSE) function otherwise.
 
         Args:
-            layer_class: Class path of a model's layer.
-            framework_attrs: Framework attributes the layer had which the graph node holds.
+            n: Node to choose distance function for.
             compute_distance_fn: An optional distance function to use globally for all nodes.
-            axis: The axis on which the operation is preformed (if specified).
             norm_mse: whether to normalize mse distance function.
 
-        Returns: A distance function between two tensors.
+        Returns: A distance function between two tensors and a axis on which the distance is computed (if exists).
         """
 
         raise NotImplemented(f'{self.__class__.__name__} have to implement the '

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/graph/base_graph.py

@@ -440,7 +440,7 @@ class Graph(nx.MultiDiGraph, GraphSearches):
 
         output_nodes = [ot.node for ot in self.get_outputs()]  # get output nodes from namedtuples
         if node_to_remove in output_nodes:  # If node is in the graph's outputs, the outputs should be updated
-            if new_graph_outputs is None:
+            if new_graph_outputs is None:  # pragma: no cover
                 Logger.critical(
                     f"{node_to_remove.name} is among the graph outputs; however, it cannot be removed without providing a new output.")  # pragma: no cover
             self.set_outputs(new_graph_outputs)
@@ -506,28 +506,6 @@ class Graph(nx.MultiDiGraph, GraphSearches):
             output_edges.sort(key=lambda e: getattr(e, sort_by_attr))
         return output_edges
 
-    def get_memory(self) -> float:
-        """
-
-        Returns: Total memory consumption of the graph in bytes.
-
-        """
-        memory = 0
-        for n in self.nodes:
-            memory += n.get_memory_bytes(self.fw_info)
-        return memory
-
-    def get_float_memory(self) -> float:
-        """
-
-        Returns: Total memory consumption of the float graph in bytes.
-
-        """
-        memory = 0
-        for n in self.nodes:
-            memory += n.get_float_memory_bytes(self.fw_info)
-        return memory
-
     def get_configurable_sorted_nodes_names(self,
                                             fw_info: FrameworkInfo,
                                             include_reused_nodes: bool = False) -> List[str]:

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/graph/base_node.py

@@ -297,18 +297,6 @@ class BaseNode:
 
         return memory
 
-    def get_float_memory_bytes(self, fw_info) -> float:
-        """
-        Compute the number of bytes the node's memory requires.
-
-        Args:
-            fw_info: Framework info to decide which attributes should be quantized.
-
-        Returns: Number of bytes the node's memory requires when in floating point (32 bit).
-
-        """
-        q_params, f_params = self.get_num_parameters(fw_info)
-        return (f_params + q_params) * FP32_BYTES_PER_PARAMETER
 
     def get_unified_weights_candidates_dict(self, fw_info) -> Dict[str, Any]:
         """
@@ -436,20 +424,6 @@ class BaseNode:
 
         return sum([np.prod([x for x in output_shape if x is not None]) for output_shape in output_shapes])
 
-    def get_total_input_params(self) -> float:
-        """
-        Calculates the total parameters in the node's input tensors.
-
-        Returns: Input size (i.e., total number of parameters).
-        """
-
-        input_shapes = self.input_shape if isinstance(self.input_shape, List) else [self.input_shape]
-
-        # remove batch size (first element) from input shape
-        input_shapes = [s[1:] for s in input_shapes]
-
-        return sum([np.prod([x for x in input_shape if x is not None]) for input_shape in input_shapes])
-
     def find_min_candidates_indices(self) -> List[int]:
         """
         Returns a list with potential minimal candidates.
@@ -644,10 +618,10 @@ class BaseNode:
         if len(simd_list) > 1:
             Logger.warning(f"More than one pruning SIMD option is available."
                            f" Min SIMD is used: {min(simd_list)}")
-        if len(simd_list) == 0:
+        if len(simd_list) == 0:  # pragma: no cover
             Logger.critical(f"No SIMD option is available for {self}")
         _simd = min(simd_list)
-        if _simd <= 0 or int(_simd) != _simd:
+        if _simd <= 0 or int(_simd) != _simd:  # pragma: no cover
             Logger.critical(f"SIMD is expected to be a non-positive integer but found: {_simd}")
         return _simd
 

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py

@@ -77,7 +77,8 @@ class SensitivityEvaluation:
         self.disable_activation_for_metric = disable_activation_for_metric
         if self.quant_config.use_hessian_based_scores:
             if not isinstance(hessian_info_service, HessianInfoService):
-                Logger.critical(f"When using Hessian-based approximations for sensitivity evaluation, a valid HessianInfoService object is required; found {type(hessian_info_service)}.")
+                Logger.critical(
+                    f"When using Hessian-based approximations for sensitivity evaluation, a valid HessianInfoService object is required; found {type(hessian_info_service)}.")
             self.hessian_info_service = hessian_info_service
 
         self.sorted_configurable_nodes_names = graph.get_configurable_sorted_nodes_names(self.fw_info)
@@ -94,7 +95,8 @@ class SensitivityEvaluation:
         self.ips_distance_fns, self.ips_axis = self._init_metric_points_lists(self.interest_points, use_normalized_mse)
 
         self.output_points = get_output_nodes_for_metric(graph)
-        self.out_ps_distance_fns, self.out_ps_axis = self._init_metric_points_lists(self.output_points, use_normalized_mse)
+        self.out_ps_distance_fns, self.out_ps_axis = self._init_metric_points_lists(self.output_points,
+                                                                                    use_normalized_mse)
 
         # Setting lists with relative position of the interest points
         # and output points in the list of all mp model activation tensors
@@ -130,7 +132,8 @@ class SensitivityEvaluation:
             self.interest_points_hessians = self._compute_hessian_based_scores()
             self.quant_config.distance_weighting_method = lambda d: self.interest_points_hessians
 
-    def _init_metric_points_lists(self, points: List[BaseNode], norm_mse: bool = False) -> Tuple[List[Callable], List[int]]:
+    def _init_metric_points_lists(self, points: List[BaseNode], norm_mse: bool = False) -> Tuple[
+        List[Callable], List[int]]:
         """
         Initiates required lists for future use when computing the sensitivity metric.
         Each point on which the metric is computed uses a dedicated distance function based on its type.
@@ -146,16 +149,12 @@ class SensitivityEvaluation:
         distance_fns_list = []
         axis_list = []
         for n in points:
-            axis = n.framework_attr.get(AXIS) if not isinstance(n, FunctionalNode) else n.op_call_kwargs.get(AXIS)
-            distance_fn = self.fw_impl.get_mp_node_distance_fn(
-                layer_class=n.layer_class,
-                framework_attrs=n.framework_attr,
-                compute_distance_fn=self.quant_config.compute_distance_fn,
-                axis=axis,
-                norm_mse=norm_mse)
+            distance_fn, axis = self.fw_impl.get_mp_node_distance_fn(n,
+                                                                     compute_distance_fn=self.quant_config.compute_distance_fn,
+                                                                     norm_mse=norm_mse)
             distance_fns_list.append(distance_fn)
             # Axis is needed only for KL Divergence calculation, otherwise we use per-tensor computation
-            axis_list.append(axis if distance_fn==compute_kl_divergence else None)
+            axis_list.append(axis if distance_fn == compute_kl_divergence else None)
         return distance_fns_list, axis_list
 
     def compute_metric(self,
@@ -300,7 +299,8 @@ class SensitivityEvaluation:
         node_name = sorted_configurable_nodes_names[node_idx_to_configure]
         layers_to_config = self.conf_node2layers.get(node_name, None)
         if layers_to_config is None:
-            Logger.critical(f"Matching layers for node {node_name} not found in the mixed precision model configuration.")  # pragma: no cover
+            Logger.critical(
+                f"Matching layers for node {node_name} not found in the mixed precision model configuration.")  # pragma: no cover
 
         for current_layer in layers_to_config:
             self.set_layer_to_bitwidth(current_layer, mp_model_configuration[node_idx_to_configure])

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/quantization/candidate_node_quantization_config.py

@@ -58,7 +58,7 @@ class CandidateNodeQuantizationConfig(BaseNodeQuantizationConfig):
         if activation_quantization_cfg is not None:
             self.activation_quantization_cfg = activation_quantization_cfg
         else:
-            if any(v is None for v in (qc, op_cfg, activation_quantization_fn, activation_quantization_params_fn)):
+            if any(v is None for v in (qc, op_cfg, activation_quantization_fn, activation_quantization_params_fn)):  # pragma: no cover
                 Logger.critical(
                     "Missing required arguments to initialize a node activation quantization configuration. "
                     "Ensure QuantizationConfig, OpQuantizationConfig, activation quantization function, "
@@ -72,7 +72,7 @@ class CandidateNodeQuantizationConfig(BaseNodeQuantizationConfig):
         if weights_quantization_cfg is not None:
             self.weights_quantization_cfg = weights_quantization_cfg
         else:
-            if any(v is None for v in (qc, op_cfg, node_attrs_list)):
+            if any(v is None for v in (qc, op_cfg, node_attrs_list)):  # pragma: no cover
                 Logger.critical("Missing required arguments to initialize a node weights quantization configuration. "
                                 "Ensure QuantizationConfig, OpQuantizationConfig, weights quantization function, "
                                 "parameters function, and weights attribute quantization config are provided.")

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/quantization/node_quantization_config.py

@@ -178,20 +178,6 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
         for param_name, param_value in activation_params.items():
             self.activation_quantization_params[param_name] = param_value
 
-    def has_activation_quantization_params(self) -> bool:
-        """
-
-        Returns: Whether NodeQuantizationConfig has a activation quantization params or not.
-
-        """
-        return len(self.activation_quantization_params) > 0
-
-    def no_quantization(self) -> bool:
-        """
-        Returns: Whether NodeQuantizationConfig does not have activation params.
-        """
-        return (not self.has_activation_quantization_params())
-
     def __eq__(self, other: Any) -> bool:
         """
         Compares the object to another object to find if they are equal.
@@ -203,7 +189,7 @@ class NodeActivationQuantizationConfig(BaseNodeQuantizationConfig):
 
         """
         if not isinstance(other, NodeActivationQuantizationConfig):
-            return False
+            return False  # pragma: no cover
 
         return self.activation_quantization_fn == other.activation_quantization_fn and \
                self.activation_quantization_params_fn == other.activation_quantization_params_fn and \
@@ -340,14 +326,6 @@ class WeightsAttrQuantizationConfig:
         else:
             self.set_weights_quantization_param({})
 
-    def has_weights_quantization_params(self) -> bool:
-        """
-
-        Returns: Whether NodeQuantizationConfig has weights quantization params or not.
-
-        """
-        return len(self.weights_quantization_params) > 0
-
     def __eq__(self, other: Any) -> bool:
         """
         Compares the object to another object to find if they are equal.
@@ -359,7 +337,7 @@ class WeightsAttrQuantizationConfig:
 
         """
         if not isinstance(other, WeightsAttrQuantizationConfig):
-            return False
+            return False  # pragma: no cover
 
         return self.weights_quantization_fn == other.weights_quantization_fn and \
                self.weights_quantization_params_fn == other.weights_quantization_params_fn and \
@@ -419,11 +397,11 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                 # In Tensorflow, the attribute name is composed of the framework attribute name and the layer name,
                 # therefore, we need to look for the attribute in the op_cfg that is contained in the node attribute's name.
                 attrs_included_in_name = {k: v for k, v in op_cfg.attr_weights_configs_mapping.items() if k in attr}
-                if len(attrs_included_in_name) > 1:
-                    Logger.error(f"Found multiple attribute in TPC OpConfig that are contained "
-                                 f"in the attribute name '{attr}'."
-                                 f"Please fix the TPC attribute names mapping such that each operator's attribute would "
-                                 f"have a unique matching name.")
+                if len(attrs_included_in_name) > 1:  # pragma: no cover
+                    Logger.critical(f"Found multiple attribute in TPC OpConfig that are contained "
+                                    f"in the attribute name '{attr}'."
+                                    f"Please fix the TPC attribute names mapping such that each operator's attribute would "
+                                    f"have a unique matching name.")
                 if len(attrs_included_in_name) == 0:
                     attr_cfg = op_cfg.default_weight_attr_config
                 else:
@@ -446,8 +424,8 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
         Returns: An attribute quantization configuration.
 
         """
-        if attr_name is None:
-            Logger.error("Got 'None' attribute name for retrieving weights attribute quantization configuration.")
+        if attr_name is None:  # pragma: no cover
+            Logger.critical("Got 'None' attribute name for retrieving weights attribute quantization configuration.")
 
         if isinstance(attr_name, int):
             # this is a positional attribute
@@ -463,8 +441,8 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                 # If no attribute with the exact name then an error would be thrown
                 attr_cfg = self.attributes_config_mapping.get(attr_name)
 
-        if attr_cfg is None:
-            Logger.error(f"Weight attribute '{attr_name}' config could not be found.")
+        if attr_cfg is None:  # pragma: no cover
+            Logger.critical(f"Weight attribute '{attr_name}' config could not be found.")
 
         return attr_cfg
 
@@ -519,8 +497,8 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                            f"{list(attrs_with_name.keys())}.")
         return attrs_with_name
 
-    def set_quant_config_attr(self, config_parameter_name: str, config_parameter_value: Any, attr_name: str = None,
-                              *args: List[Any], **kwargs: Dict[str, Any]):
+    def set_quant_config_attr(self, config_parameter_name: str, config_parameter_value: Any,
+                              attr_name: Union[str, int] = None, *args: List[Any], **kwargs: Dict[str, Any]):
         """
         This method overrides the parent class set_quant_config_attr to enable setting a specific weights
         attribute config parameter.
@@ -546,8 +524,8 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
                 else:
                     Logger.warning(f"Parameter {config_parameter_name} could not be found in the node quantization config of "
                                    f"weights attribute {attr_name} and was not updated!")
-            else:
-                Logger.error(f"Weights attribute {attr_name} could not be found to set parameter {config_parameter_name}.")
+            else:  # pragma: no cover
+                Logger.critical(f"Weights attribute {attr_name} could not be found to set parameter {config_parameter_name}.")
 
     def __eq__(self, other: Any) -> bool:
         """
@@ -560,7 +538,7 @@ class NodeWeightsQuantizationConfig(BaseNodeQuantizationConfig):
 
         """
         if not isinstance(other, NodeWeightsQuantizationConfig):
-            return False
+            return False  # pragma: no cover
 
         return self.min_threshold == other.min_threshold and \
             self.simd_size == other.simd_size and \

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_weights_computation.py

@@ -22,10 +22,7 @@ from model_compression_toolkit.defaultdict import DefaultDict
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.quantization.node_quantization_config import NodeWeightsQuantizationConfig, \
     WeightsAttrQuantizationConfig
-
-# If the quantization config does not contain kernel channel mapping or the weights
-# quantization is not per-channel, we use a dummy channel mapping.
-dummy_channel_mapping = DefaultDict(default_value=(None, None))
+from model_compression_toolkit.logger import Logger
 
 
 def get_weights_qparams(weights_attr_values: np.ndarray,
@@ -64,29 +61,10 @@ def get_weights_qparams(weights_attr_values: np.ndarray,
             node=node,
             hessian_info_service=hessian_info_service,
             num_hessian_samples=num_hessian_samples)
-    else:
+    else:  # pragma: no cover
+        Logger.error(f"Requested weights quantization parameters computation for node {node.name} without providing a "
+                     f"weights_quantization_params_fn."
+                     f"Returning an empty dictionary since no quantization parameters were computed.")
         weights_params = {}
 
     return weights_params, output_channels_axis
-
-
-def _get_kernel_channels_mapping(fw_info:FrameworkInfo,
-                                use_dummy: bool) -> DefaultDict:
-    """
-    Get a kernel channel mapping from the framework info, or use dummy mapping (which returns a
-    tuple of Nones) if use_use_dummy is True.
-
-    Args:
-        fw_info: Framework info which contains a kernel channels mapping.
-        use_dummy: Whether to use a dummy mapping or not.
-
-    Returns:
-        Kernel channels mapping.
-    """
-
-    # Set a kernel channels mapping
-    if use_dummy:  # If kernel mapping is missing, we use a dummy channels mapping
-        kernel_channels_mapping = dummy_channel_mapping
-    else:
-        kernel_channels_mapping = fw_info.kernel_channels_mapping
-    return kernel_channels_mapping

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/common/quantization/quantizers/quantizers_helpers.py

@@ -60,33 +60,6 @@ def calculate_delta(threshold: np.ndarray,
     return threshold / (2 ** (n_bits - int(signed)))
 
 
-def calculate_min_max_values(threshold: np.ndarray,
-                             n_bits: int = 8,
-                             signed: bool = False) -> Tuple[np.ndarray, np.ndarray]:
-    """
-    Compute the min/max values of a quantization range according to the threshold,
-    number of bits and whether its signed or unsigned.
-
-    Args:
-        threshold: Threshold of quantization range to compute its min/max values.
-        n_bits: Number of bits used in the quantization.
-        signed: Whether the quantization range is signed or not.
-
-    Returns:
-        Min/max values of quantization range.
-    """
-
-    delta = calculate_delta(threshold,
-                            n_bits=n_bits,
-                            signed=signed)
-
-    # If unsigned: min=0, otherwise its -threshold
-    min_value = int(signed) * -threshold
-    max_value = threshold - delta
-
-    return min_value, max_value
-
-
 def quantize_tensor(tensor_data: np.ndarray,
                     threshold: np.ndarray,
                     n_bits: int,
@@ -238,7 +211,7 @@ def get_tensor_max(tensor_data: np.ndarray,
 
     """
     if n_bits < 1:
-        Logger.critical(f"Parameter n_bits must be positive; however 'n_bits'={n_bits} was provided.")
+        Logger.critical(f"Parameter n_bits must be positive; however 'n_bits'={n_bits} was provided.")  # pragma: no cover
     if is_uniform_quantization:
         expansion_factor = 1.0
     elif n_bits == 1:
@@ -337,40 +310,3 @@ def get_output_shape(tensor_shape, channel_axis):
 
     """
     return [-1 if i is channel_axis else 1 for i in range(len(tensor_shape))]
-
-
-def get_range_bounds(tensor_min, tensor_max):
-    """
-    Gets bounds on the quantization range limits for the minimization process.
-    Calculates the bounds in a way that would leave a gap between the possible optimized values
-    and the tensor min-max values.
-
-    Args:
-        tensor_min: min value of a tensor.
-        tensor_max: max value of a tensor.
-
-    Returns: An array with (lbound, ubound) pairs on the quantization range limit values.
-
-    """
-    # choosing bounds that have some gap from the original tensor min/max values.
-    l_bound = tensor_min / 2 if tensor_min > 0 else tensor_min * 2
-    u_bound = tensor_max * 2 if tensor_max > 0 else tensor_min / 2
-    return [(l_bound, u_bound), (l_bound, u_bound)]
-
-
-def get_threshold_bounds(min_threshold, max_threshold):
-    """
-    Gets bounds on the threshold for the minimization process.
-    Calculates the bounds in a way that would leave a gap between the possible optimized threshold
-    and the tensor max values. We use min_threshold as lower-bound to prevent the selected threshold
-    from being zero or negative.
-
-    Args:
-        min_threshold: minimal threshold to use if threshold is too small (not used for this method).
-        max_threshold: maximal threshold to be used in quantization.
-
-    Returns: An array with a pair of (lbound, ubound) on the quantization threshold limit values.
-
-    """
-    max_threshold = max(min_threshold, max_threshold)
-    return [(min_threshold, 2 * max_threshold)]

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/keras/back2framework/factory_model_builder.py

@@ -37,10 +37,10 @@ def get_keras_model_builder(mode: ModelBuilderMode) -> type:
         Keras model builder for the given mode.
     """
 
-    if not isinstance(mode, ModelBuilderMode):
+    if not isinstance(mode, ModelBuilderMode):  # pragma: no cover
         Logger.critical(f"Expected a ModelBuilderMode type for 'mode', but received {type(mode)} instead.")
-    if mode is None:
+    if mode is None:  # pragma: no cover
         Logger.critical(f"get_keras_model_builder received 'mode' is None")
-    if mode not in keras_model_builders.keys():
+    if mode not in keras_model_builders.keys():  # pragma: no cover
         Logger.critical(f"'mode' {mode} is not recognized in the Keras model builders factory.")
     return keras_model_builders.get(mode)

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/keras/keras_implementation.py

@@ -21,6 +21,7 @@ from mct_quantizers import KerasQuantizationWrapper, KerasActivationQuantization
 from tensorflow.keras.models import Model
 
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
+from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianMode, HessianInfoService
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.remove_identity import RemoveIdentity
 from model_compression_toolkit.core.keras.hessian.activation_hessian_scores_calculator_keras import \
@@ -421,44 +422,47 @@ class KerasImplementation(FrameworkImplementation):
 
         return False
 
-    def get_mp_node_distance_fn(self, layer_class: type,
-                             framework_attrs: Dict[str, Any],
-                             compute_distance_fn: Callable = None,
-                             axis: int = None,
-                             norm_mse: bool = False) -> Callable:
+    def get_mp_node_distance_fn(self, n: BaseNode,
+                                compute_distance_fn: Callable = None,
+                                norm_mse: bool = False) -> Tuple[Callable, int]:
         """
         A mapping between layers' types and a distance function for computing the distance between
         two tensors in mixed precision (for loss computation purposes). Returns a specific function if node of specific types is
         given, or a default (normalized MSE) function otherwise.
 
         Args:
-            layer_class: Class path of a model's layer.
-            framework_attrs: Framework attributes the layer had which the graph node holds.
+            n: Node to choose distance function for.
             compute_distance_fn: An optional distance function to use globally for all nodes.
-            axis: The axis on which the operation is preformed (if specified).
             norm_mse: whether to normalize mse distance function.
 
-        Returns: A distance function between two tensors.
+        Returns: A distance function between two tensors and a axis on which the distance is computed (if exists).
         """
 
+        axis = n.framework_attr.get(keras_constants.AXIS) \
+            if not isinstance(n, FunctionalNode) else n.op_call_kwargs.get(keras_constants.AXIS)
+
+        layer_class = n.layer_class
+        framework_attrs = n.framework_attr
+
         if compute_distance_fn is not None:
-            return compute_distance_fn
+            return compute_distance_fn, axis
 
         if layer_class == Activation:
             node_type_name = framework_attrs[ACTIVATION]
             if node_type_name == SOFTMAX and axis is not None:
-                return compute_kl_divergence
+                return compute_kl_divergence, axis
             elif node_type_name == SIGMOID:
-                return compute_cs
+                return compute_cs, axis
         elif axis is not None and (layer_class == tf.nn.softmax or layer_class == tf.keras.layers.Softmax
                                    or (layer_class == TFOpLambda and
                                        SOFTMAX in framework_attrs[keras_constants.FUNCTION])):
-            return compute_kl_divergence
-        elif layer_class == tf.nn.sigmoid:
-            return compute_cs
+            return compute_kl_divergence, axis
+        elif layer_class == tf.nn.sigmoid or (layer_class == TFOpLambda and
+                                              SIGMOID in framework_attrs[keras_constants.FUNCTION]):
+            return compute_cs, axis
         elif layer_class == Dense:
-            return compute_cs
-        return partial(compute_mse, norm=norm_mse)
+            return compute_cs, axis
+        return partial(compute_mse, norm=norm_mse), axis
 
     def get_hessian_scores_calculator(self,
                                       graph: Graph,

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/keras/tf_tensor_numpy.py

@@ -36,7 +36,7 @@ def to_tf_tensor(tensor):
         return (to_tf_tensor(t) for t in tensor)
     elif isinstance(tensor, np.ndarray):
         return tf.convert_to_tensor(tensor.astype(np.float32))
-    else:
+    else:  # pragma: no cover
         Logger.critical(f'Unsupported type for conversion to TF tensor: {type(tensor)}.')
 
 
@@ -69,5 +69,5 @@ def tf_tensor_to_numpy(tensor: Union[List, Tuple, np.ndarray, tf.Tensor, float],
         return np.array([np_tensor]) if np.isscalar(np_tensor) else np_tensor
     elif isinstance(tensor, float):
         return np.array([tensor])
-    else:
+    else:  # pragma: no cover
         Logger.critical(f'Unsupported type for conversion to Numpy array: {type(tensor)}.')

{mct-nightly-2.1.0.20240731.414 → mct-nightly-2.1.0.20240801.515}/model_compression_toolkit/core/pytorch/back2framework/factory_model_builder.py

@@ -37,10 +37,10 @@ def get_pytorch_model_builder(mode: ModelBuilderMode) -> type:
         PyTorch model builder for the given mode.
     """
 
-    if not isinstance(mode, ModelBuilderMode):
+    if not isinstance(mode, ModelBuilderMode):  # pragma: no cover
         Logger.critical(f"Expected a ModelBuilderMode type for 'mode' parameter; received {type(mode)} instead.")
-    if mode is None:
+    if mode is None:  # pragma: no cover
         Logger.critical(f"Received 'mode' parameter is None.")
-    if mode not in pytorch_model_builders.keys():
+    if mode not in pytorch_model_builders.keys():  # pragma: no cover
         Logger.critical(f"'mode' parameter {mode} is not supported by the PyTorch model builders factory.")
     return pytorch_model_builders.get(mode)