mct-nightly 2.1.0.20240609.524__tar.gz → 2.1.0.20240611.428__tar.gz

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/PKG-INFO

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

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/mct_nightly.egg-info/PKG-INFO

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

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/mct_nightly.egg-info/SOURCES.txt

@@ -462,6 +462,14 @@ model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tp_model.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_keras.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tpc_pytorch.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/__init__.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/tp_model.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/tpc_keras.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3/tpc_pytorch.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/__init__.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tp_model.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_keras.py
+model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v3_lut/tpc_pytorch.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/__init__.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/target_platform_capabilities.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/latest/__init__.py

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/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.20240609.000524"
+__version__ = "2.1.0.20240611.000428"

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/graph/base_node.py

@@ -240,10 +240,7 @@ class BaseNode:
                                   if isinstance(pos, int)):
             if pos > len(input_tensors):
                 Logger.critical("The positional weight index cannot exceed the number of input tensors to the node.")  # pragma: no cover
-            # Insert only positional weights that are not subject to quantization. If the positional weight is
-            # subject to quantization, the quantization wrapper inserts the positional weight into the node.
-            if not self.is_weights_quantization_enabled(pos):
-                input_tensors.insert(pos, weight)
+            input_tensors.insert(pos, weight)
 
         return input_tensors
 

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/hessian/hessian_info_service.py

@@ -388,4 +388,4 @@ class HessianInfoService:
         Returns: A list with split batch into individual results.
 
         """
-        return [t[i] for i in range(t.shape[0])]
+        return [t[i:i+1, :] for i in range(t.shape[0])]

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/pruning/importance_metrics/lfh_importance_metric.py

@@ -135,7 +135,7 @@ class LFHImportanceMetric(BaseImportanceMetric):
             nodes_scores.append(_scores_for_node)
 
         # Average and map scores to nodes.
-        self._entry_node_to_hessian_score = {node: np.mean(scores[0], axis=0) for node, scores in zip(entry_nodes, nodes_scores)}
+        self._entry_node_to_hessian_score = {node: np.mean(scores[0], axis=0).flatten() for node, scores in zip(entry_nodes, nodes_scores)}
 
         self._entry_node_count_oc_nparams = self._count_oc_nparams(entry_nodes=entry_nodes)
         _entry_node_l2_oc_norm = self._get_squaredl2norm(entry_nodes=entry_nodes)

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/quantization/node_quantization_config.py

@@ -326,13 +326,17 @@ class WeightsAttrQuantizationConfig:
 
         """
         assert self.enable_weights_quantization
+        assert not (self.weights_per_channel_threshold and self.weights_channels_axis is None), \
+            "Trying to calculate threshold per channel, channel axis in None."
         if self.weights_quantization_params_fn is not None:
-            self.set_weights_quantization_param(self.weights_quantization_params_fn(tensor_data,
-                                                                                    p=self.l_p_value,
-                                                                                    n_bits=self.weights_n_bits,
-                                                                                    per_channel=self.weights_per_channel_threshold and self.weights_channels_axis is not None,
-                                                                                    channel_axis=self.weights_channels_axis[0],  # output channel axis
-                                                                                    min_threshold=min_threshold))
+            self.set_weights_quantization_param(
+                self.weights_quantization_params_fn(tensor_data,
+                                                    p=self.l_p_value,
+                                                    n_bits=self.weights_n_bits,
+                                                    per_channel=self.weights_per_channel_threshold and self.weights_channels_axis is not None,
+                                                    channel_axis=self.weights_channels_axis[0],  # output channel axis
+                                                    min_threshold=min_threshold)[0]  # Take only first output, the q-params, as axis is already chosen.
+            )
         else:
             self.set_weights_quantization_param({})
 

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/quantization/quantization_params_generation/lut_kmeans_params.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Dict
+from typing import Dict, Tuple
 import numpy as np
 from sklearn.cluster import KMeans
 
@@ -42,7 +42,8 @@ def lut_kmeans_tensor(tensor_data: np.ndarray,
                       is_symmetric: bool = False,
                       node=None,
                       hessian_info_service: HessianInfoService = None,
-                      num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Dict:
+                      num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES,
+                      ) -> Tuple[Dict[str, np.ndarray], int]:
     """
     The quantizer first finds the closest max value per channel of tensor_data.
     Now, we divide tensor_data with the threshold vector per channel. In addition, we scale the result to the range
@@ -70,27 +71,34 @@ def lut_kmeans_tensor(tensor_data: np.ndarray,
     if n_bits >= LUT_VALUES_BITWIDTH:
         Logger.critical(f'Look-Up-Table (LUT) bit configuration exceeds maximum: {n_bits} bits provided, must be less than {LUT_VALUES_BITWIDTH} bits.')  # pragma: no cover
     # TODO: need to set this externally
+    n_data_points = len(np.unique(tensor_data.flatten()))
     if len(np.unique(tensor_data.flatten())) < 2 ** n_bits:
-        n_clusters = len(np.unique(tensor_data.flatten()))
+        n_clusters = n_data_points
     else:
         n_clusters = 2 ** n_bits
     kmeans = KMeans(n_clusters=n_clusters, n_init=10)
 
     threshold_selection_tensor = symmetric_selection_tensor if is_symmetric else power_of_two_selection_tensor
-    thresholds_per_channel = threshold_selection_tensor(tensor_data, p, n_bits, per_channel,
-                                                        channel_axis, n_iter, min_threshold,
-                                                        qc.QuantizationErrorMethod.NOCLIPPING)[THRESHOLD]
+
+    _params, channel_axis = threshold_selection_tensor(tensor_data, p, n_bits, per_channel,
+                                                       channel_axis, n_iter, min_threshold,
+                                                       qc.QuantizationErrorMethod.NOCLIPPING)
+    thresholds_per_channel = _params[THRESHOLD]
 
     tensor_for_kmeans = int_quantization_with_threshold(tensor_data, thresholds_per_channel, LUT_VALUES_BITWIDTH)
     kmeans.fit(tensor_for_kmeans.reshape(-1, 1))
 
     # Add 0 to the LUT
     cc = np.round(kmeans.cluster_centers_)
+    if n_data_points < 2 ** n_bits and np.all(cc != 0):
+        # In case there are fewer data points than potential clusters, we can add the cluster 0.0
+        # to the original clusters array to improve quantization (i.e. no need to zero one of the clusters).
+        cc = np.concatenate([np.zeros([1, 1], dtype=cc.dtype), cc])
     closest2zero_idx = (np.abs(cc - 0)).argmin()
     cc[closest2zero_idx] = 0.0
 
     return {LUT_VALUES: cc,
-            SCALE_PER_CHANNEL: thresholds_per_channel}
+            SCALE_PER_CHANNEL: thresholds_per_channel}, channel_axis
 
 
 def lut_kmeans_histogram(bins: np.ndarray,

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/quantization/quantization_params_generation/power_of_two_selection.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 import numpy as np
+from typing import Union, Tuple, Dict
 
 import model_compression_toolkit.core.common.quantization.quantization_config as qc
 from model_compression_toolkit.constants import MIN_THRESHOLD, THRESHOLD, NUM_QPARAM_HESSIAN_SAMPLES
@@ -23,20 +24,22 @@ from model_compression_toolkit.core.common.quantization.quantizers.quantizers_he
 from model_compression_toolkit.core.common.quantization.quantization_params_generation.error_functions import \
     get_threshold_selection_tensor_error_function, get_threshold_selection_histogram_error_function
 from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from model_compression_toolkit.core.common.similarity_analyzer import compute_mse
+from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import quantize_tensor
 
 
 def power_of_two_selection_tensor(tensor_data: np.ndarray,
                                   p: int,
                                   n_bits: int,
                                   per_channel: bool = False,
-                                  channel_axis: int = 1,
+                                  channel_axis: Union[int, None] = 1,
                                   n_iter: int = 10,
                                   min_threshold: float = MIN_THRESHOLD,
                                   quant_error_method: qc.QuantizationErrorMethod = qc.QuantizationErrorMethod.MSE,
                                   node=None,
                                   hessian_info_service: HessianInfoService = None,
                                   num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES,
-                                  ) -> dict:
+                                  ) -> Tuple[Dict[str, np.ndarray], int]:
     """
     Compute the power of two threshold based on the provided QuantizationErrorMethod to quantize the tensor.
     Different search is applied, depends on the value of the selected QuantizationErrorMethod.
@@ -46,7 +49,7 @@ def power_of_two_selection_tensor(tensor_data: np.ndarray,
         p: p-norm to use for the Lp-norm distance.
         n_bits: Number of bits to quantize the tensor.
         per_channel: Whether the quantization should be per-channel or not.
-        channel_axis: Output channel index.
+        channel_axis: Output channel index. if None, search for best axis.
         n_iter: Number of iterations to search for the optimal threshold (not used for this method).
         min_threshold: Minimal threshold to use if threshold is too small (not used for this method).
         quant_error_method: an error function to optimize the parameters' selection accordingly.
@@ -56,11 +59,24 @@ def power_of_two_selection_tensor(tensor_data: np.ndarray,
 
     Returns:
         Power of two threshold to quantize the tensor in a power of 2 manner.
+        Selected quantization channel axis.
     """
 
     if quant_error_method == qc.QuantizationErrorMethod.NOCLIPPING:
-        tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits)
-        threshold = max_power_of_two(tensor_max, min_threshold)
+        if channel_axis is None and per_channel:
+            total_error_list = []
+            th_list = []
+            for _axis in range(len(tensor_data.shape)):
+                tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits)
+                threshold = max_power_of_two(tensor_max, min_threshold)
+                q_tensor_data = quantize_tensor(tensor_data, threshold, n_bits, True)
+                total_error_list.append(compute_mse(tensor_data, q_tensor_data, norm=True))
+                th_list.append(threshold)
+            channel_axis = np.argmin(total_error_list)
+            threshold = th_list[channel_axis]
+        else:
+            tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits)
+            threshold = max_power_of_two(tensor_max, min_threshold)
     else:
         signed = True  # weights are always signed
         axis = -1 if per_channel else None
@@ -69,15 +85,15 @@ def power_of_two_selection_tensor(tensor_data: np.ndarray,
                                                                        n_bits=n_bits, signed=signed, node=node,
                                                                        hessian_info_service=hessian_info_service,
                                                                        num_hessian_samples=num_hessian_samples)
-        threshold = qparams_selection_tensor_search(error_function,
-                                                    tensor_data,
-                                                    n_bits,
-                                                    per_channel=per_channel,
-                                                    channel_axis=channel_axis,
-                                                    n_iter=n_iter,
-                                                    min_threshold=min_threshold,
-                                                    signed=signed)
-    return {THRESHOLD: threshold}
+        threshold, channel_axis = qparams_selection_tensor_search(error_function,
+                                                                  tensor_data,
+                                                                  n_bits,
+                                                                  per_channel=per_channel,
+                                                                  channel_axis=channel_axis,
+                                                                  n_iter=n_iter,
+                                                                  min_threshold=min_threshold,
+                                                                  signed=signed)
+    return {THRESHOLD: threshold}, channel_axis
 
 
 def power_of_two_selection_histogram(bins: np.ndarray,

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_computation.py

@@ -84,13 +84,14 @@ def calculate_quantization_params(graph: Graph,
                             mod_attr_cfg = copy.deepcopy(attr_cfg)
                             mod_attr_cfg.weights_error_method = QuantizationErrorMethod.MSE
 
-                    weights_params = get_weights_qparams(n.get_weights_by_keys(attr),
-                                                         candidate_qc.weights_quantization_cfg,
-                                                         mod_attr_cfg,
-                                                         output_channels_axis,
-                                                         node=n,
-                                                         hessian_info_service=hessian_info_service,
-                                                         num_hessian_samples=num_hessian_samples)
+                    weights_params, output_channels_axis = get_weights_qparams(n.get_weights_by_keys(attr),
+                                                                               candidate_qc.weights_quantization_cfg,
+                                                                               mod_attr_cfg,
+                                                                               output_channels_axis,
+                                                                               node=n,
+                                                                               hessian_info_service=hessian_info_service,
+                                                                               num_hessian_samples=num_hessian_samples)
+                    attr_cfg.weights_channels_axis = (output_channels_axis, attr_cfg.weights_channels_axis[1])
                     attr_cfg.set_weights_quantization_param(weights_params)
 
             if n.is_activation_quantization_enabled():

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_search.py

@@ -27,7 +27,7 @@ from model_compression_toolkit.constants import MIN_THRESHOLD, DEFAULT_TOL, DEFA
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import quantize_tensor, \
     reshape_tensor_for_per_channel_search, uniform_quantize_tensor, get_output_shape
 from model_compression_toolkit.core.common.quantization.quantizers.quantizers_helpers import max_power_of_two, \
-    get_tensor_max
+    get_tensor_max, get_tensor_min
 
 
 def qparams_selection_tensor_search(error_function: Callable,
@@ -56,41 +56,49 @@ def qparams_selection_tensor_search(error_function: Callable,
         signed: a flag whether the tensor is signed.
 
     Returns:
-        Optimal constrained threshold to quantize the tensor.
+        Optimal constrained threshold to quantize the tensor, and best channel axis if input channel_axis was None,
+        else return the input channel axis.
 
     """
 
-    output_shape = get_output_shape(tensor_data.shape, channel_axis)
-
-    # First threshold to check is the constrained threshold based on the tensor's maximal value.
-    tensor_max = get_tensor_max(tensor_data, per_channel, channel_axis, n_bits)
-    threshold = 2 * max_power_of_two(tensor_max, min_threshold)
+    search_axes = range(len(tensor_data.shape)) if channel_axis is None and per_channel else [channel_axis]
+    total_error_list = []
+    th_list = []
+    for _axis in search_axes:
+        output_shape = get_output_shape(tensor_data.shape, _axis)
 
-    # If the threshold is computed per-channel, we rearrange the tensor such that each sub-tensor
-    # is flattened, and we iterate over each one of them when searching for the threshold.
-    if per_channel:
-        tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, channel_axis)
+        # First threshold to check is the constrained threshold based on the tensor's maximal value.
+        tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits)
+        threshold = 2 * max_power_of_two(tensor_max, min_threshold)
 
-    error_list = []  # init an empty error list
-    # On each iteration a new constrained threshold which equal to half of the previous tested threshold
-    # is used for quantizing the tensor and computing the error. The error is appended to an error list, which
-    # eventually used to select the threshold with the minimal error.
-    for i in range(n_iter):
+        # Rearrange the tensor such that each sub-tensor is flattened, and we iterate over each
+        # one of them when searching for the threshold.
         if per_channel:
-            threshold_hat = (threshold / (2 ** i)).reshape([-1, 1])
-            qt = quantize_tensor(tensor_data_r, threshold_hat, n_bits, signed)
-            per_channel_error = _error_function_wrapper(error_function, tensor_data_r, qt, threshold_hat)
-
-            error_list.append(per_channel_error)
-        else:  # quantize per-tensor
-            qt = quantize_tensor(tensor_data, threshold / (2 ** i), n_bits, signed)
-            error = error_function(qt, tensor_data, threshold=threshold / (2 ** i))
-            error_list.append(error)
-
-    # Take the index of the minimal error, and use it compute the threshold which yielded it.
-    i = np.argmin(np.stack(error_list, axis=-1), axis=-1)
-
-    return np.maximum(np.reshape(threshold.flatten() / np.power(2, i), output_shape), min_threshold)
+            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, _axis)
+
+        error_list = []  # init an empty error list
+        # On each iteration a new constrained threshold which equal to half of the previous tested threshold
+        # is used for quantizing the tensor and computing the error. The error is appended to an error list, which
+        # eventually used to select the threshold with the minimal error.
+        for i in range(n_iter):
+            if per_channel:
+                threshold_hat = (threshold / (2 ** i)).reshape([-1, 1])
+                qt = quantize_tensor(tensor_data_r, threshold_hat, n_bits, signed)
+                per_channel_error = _error_function_wrapper(error_function, tensor_data_r, qt, threshold_hat)
+                error_list.append(per_channel_error)
+            else:  # quantize per-tensor
+                qt = quantize_tensor(tensor_data, threshold / (2 ** i), n_bits, signed)
+                error = error_function(qt, tensor_data, threshold=threshold / (2 ** i))
+                error_list.append(error)
+
+        # Take the index of the minimal error, and use it compute the threshold which yielded it.
+        err_mat = np.stack(error_list, axis=-1)
+        i = np.argmin(err_mat, axis=-1)
+        th_list.append(np.maximum(np.reshape(threshold.flatten() / np.power(2, i), output_shape), min_threshold))
+        total_error_list.append(err_mat.min(axis=-1).mean())
+
+    best_axis_index = np.argmin(total_error_list)
+    return th_list[best_axis_index], search_axes[best_axis_index]
 
 
 def qparams_selection_histogram_search(error_function: Callable,
@@ -390,13 +398,12 @@ def search_dynamic_range(base_range: np.ndarray, x: np.ndarray, scalers: np.ndar
 
 def qparams_symmetric_selection_tensor_search(error_function: Callable,
                                               tensor_data: np.ndarray,
-                                              tensor_max: np.ndarray,
                                               n_bits: int,
                                               per_channel: bool = False,
                                               channel_axis: int = 1,
                                               n_iter: int = SYMMETRIC_TENSOR_PER_CHANNEL_N_ITER,
                                               min_threshold=MIN_THRESHOLD,
-                                              signed: bool = True) -> Any:
+                                              signed: bool = True) -> Tuple[np.ndarray, int]:
     """
     Search for optimal threshold (per-channel or per-tensor) for symmetric quantization of a tensor,
     using the iterative optimizer method.
@@ -404,7 +411,6 @@ def qparams_symmetric_selection_tensor_search(error_function: Callable,
     Args:
         error_function: Function to compute the error between the original and quantized tensors.
         tensor_data: Numpy array with tensor's content.
-        tensor_max: The max value of the tensor.
         n_bits: Number of bits to quantize the tensor.
         per_channel: Whether the tensor should be quantized per-channel or per-tensor.
         channel_axis: Index of output channels dimension.
@@ -417,46 +423,55 @@ def qparams_symmetric_selection_tensor_search(error_function: Callable,
 
     """
 
-    output_shape = get_output_shape(tensor_data.shape, channel_axis)
+    search_axes = range(len(tensor_data.shape)) if channel_axis is None and per_channel else [channel_axis]
+    total_error_list = []
+    th_list = []
+    for _axis in search_axes:
+        tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits)
+        output_shape = get_output_shape(tensor_data.shape, _axis)
 
-    # If the threshold is computed per-channel, we rearrange the tensor such that each sub-tensor
-    # is flattened, and we iterate over each one of them when searching for the threshold.
-    if per_channel:
-        tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, channel_axis)
-        max_tensor = np.maximum(min_threshold, tensor_max)
-        res = qparams_symmetric_iterative_minimization(x0=max_tensor,
-                                                       x=tensor_data_r,
-                                                       loss_fn=error_function,  # gets float_tensor, fxp_tensor, threshold
-                                                       n_bits=n_bits,
-                                                       signed=signed,
-                                                       n_intervals=SYMMETRIC_TENSOR_PER_CHANNEL_N_INTERVALS,
-                                                       n_iter=SYMMETRIC_TENSOR_PER_CHANNEL_N_ITER,
-                                                       dec_freq=SYMMETRIC_TENSOR_PER_CHANNEL_DEC_FREQ,
-                                                       per_channel=True)
-        return np.reshape(np.maximum(min_threshold, res['param']), output_shape)
-    else:
-        # quantize per-tensor
-        res = qparams_symmetric_iterative_minimization(x0=get_init_threshold(min_threshold, tensor_max),
-                                                       x=tensor_data,
-                                                       loss_fn=error_function,
-                                                       n_bits=n_bits,
-                                                       signed=signed,
-                                                       n_intervals=SYMMETRIC_TENSOR_N_INTERVALS,
-                                                       n_iter=SYMMETRIC_TENSOR_N_ITER,
-                                                       dec_freq=SYMMETRIC_TENSOR_DEC_FREQ,
-                                                       per_channel=False)
-
-        return max(min_threshold, res['param'])
+        if per_channel:
+            # Rearrange the tensor such that each sub-tensor is flattened, and we iterate
+            # over each one of them when searching for the threshold.
+            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, _axis)
+            max_tensor = np.maximum(min_threshold, tensor_max)
+            res = qparams_symmetric_iterative_minimization(x0=max_tensor,
+                                                           x=tensor_data_r,
+                                                           loss_fn=error_function,  # gets float_tensor, fxp_tensor, threshold
+                                                           n_bits=n_bits,
+                                                           signed=signed,
+                                                           n_intervals=SYMMETRIC_TENSOR_PER_CHANNEL_N_INTERVALS,
+                                                           n_iter=SYMMETRIC_TENSOR_PER_CHANNEL_N_ITER,
+                                                           dec_freq=SYMMETRIC_TENSOR_PER_CHANNEL_DEC_FREQ,
+                                                           per_channel=True)
+            th = np.reshape(np.maximum(min_threshold, res['param']), output_shape)
+        else:
+            # quantize per-tensor
+            res = qparams_symmetric_iterative_minimization(x0=get_init_threshold(min_threshold, tensor_max),
+                                                           x=tensor_data,
+                                                           loss_fn=error_function,
+                                                           n_bits=n_bits,
+                                                           signed=signed,
+                                                           n_intervals=SYMMETRIC_TENSOR_N_INTERVALS,
+                                                           n_iter=SYMMETRIC_TENSOR_N_ITER,
+                                                           dec_freq=SYMMETRIC_TENSOR_DEC_FREQ,
+                                                           per_channel=False)
+            th = max(min_threshold, res['param'])
+
+        total_error_list.append(res['loss'].mean())
+        th_list.append(th)
+
+    best_axis_index = np.argmin(total_error_list)
+    return th_list[best_axis_index], search_axes[best_axis_index]
 
 
 def qparams_uniform_selection_tensor_search(error_function: Callable,
                                             tensor_data: np.ndarray,
-                                            tensor_min: np.ndarray,
-                                            tensor_max: np.ndarray,
                                             n_bits: int,
                                             per_channel: bool = False,
                                             channel_axis: int = 1,
-                                            n_iter: int = UNIFORM_TENSOR_PER_CHANNEL_N_ITER) -> Any:
+                                            n_iter: int = UNIFORM_TENSOR_PER_CHANNEL_N_ITER,
+                                            ) -> Tuple[Tuple[np.ndarray, np.ndarray], int]:
     """
     Search for optimal quantization range (per-channel or per-tensor) for uniform quantization of a tensor,
     using the iterative optimizer method and built-in scale factors
@@ -465,8 +480,6 @@ def qparams_uniform_selection_tensor_search(error_function: Callable,
     Args:
         error_function: Function to compute the error between the original and quantized tensors.
         tensor_data: Numpy array with tensor's content.
-        tensor_min: The min value of the tensor.
-        tensor_max: The max value of the tensor.
         n_bits: Number of bits to quantize the tensor.
         per_channel: Whether the tensor should be quantized per-channel or per-tensor.
         channel_axis: Index of output channels dimension.
@@ -477,17 +490,22 @@ def qparams_uniform_selection_tensor_search(error_function: Callable,
 
     """
 
-    output_shape = get_output_shape(tensor_data.shape, channel_axis)
+    search_axes = range(len(tensor_data.shape)) if channel_axis is None and per_channel else [channel_axis]
+    total_error_list = []
+    th_list = []
+    for _axis in search_axes:
+        tensor_min = get_tensor_min(tensor_data, per_channel, _axis)
+        tensor_max = get_tensor_max(tensor_data, per_channel, _axis, n_bits, is_uniform_quantization=True)
+        output_shape = get_output_shape(tensor_data.shape, _axis)
 
-    alpha = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
-    beta = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
-    scalers = np.asarray(list(itertools.product(alpha, beta)))
+        alpha = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
+        beta = np.linspace(BOTTOM_FACTOR, UPPER_FACTOR, UNIFORM_TENSOR_N_SAMPLES)
+        scalers = np.asarray(list(itertools.product(alpha, beta)))
 
-    # If the threshold is computed per-channel, we rearrange the tensor such that each sub-tensor
-    # is flattened, and we iterate over each one of them when searching for the threshold.
-    if per_channel:
+        # Rearrange the tensor such that each sub-tensor is flattened, and we iterate over
+        # each one of them when searching for the threshold.
         if per_channel:
-            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, channel_axis)
+            tensor_data_r = reshape_tensor_for_per_channel_search(tensor_data, _axis)
             tensor_min_max = np.column_stack([tensor_min.flatten(), tensor_max.flatten()])
             res = iterative_uniform_dynamic_range_search(x0=tensor_min_max,
                                                          x=tensor_data_r,
@@ -496,18 +514,21 @@ def qparams_uniform_selection_tensor_search(error_function: Callable,
                                                          n_bits=n_bits,
                                                          n_iter=UNIFORM_TENSOR_PER_CHANNEL_N_ITER,
                                                          per_channel=True)
-            return np.reshape(res['param'][:, 0], output_shape), np.reshape(res['param'][:, 1], output_shape)
-    else:
-        # quantize per-tensor
-        pass
-        res = iterative_uniform_dynamic_range_search(x0=np.array([tensor_min, tensor_max]),
-                                                     x=tensor_data,
-                                                     scalers=scalers,
-                                                     loss_fn=error_function,
-                                                     n_bits=n_bits,
-                                                     n_iter=UNIFORM_TENSOR_N_ITER,
-                                                     per_channel=False)
-        return res['param']
+            th_list.append((np.reshape(res['param'][:, 0], output_shape), np.reshape(res['param'][:, 1], output_shape)))
+        else:
+            # quantize per-tensor
+            res = iterative_uniform_dynamic_range_search(x0=np.array([tensor_min, tensor_max]),
+                                                         x=tensor_data,
+                                                         scalers=scalers,
+                                                         loss_fn=error_function,
+                                                         n_bits=n_bits,
+                                                         n_iter=UNIFORM_TENSOR_N_ITER,
+                                                         per_channel=False)
+            th_list.append(tuple(np.split(res['param'], 2)))
+        total_error_list.append(res['loss'].mean())
+
+    best_axis_index = np.argmin(total_error_list)
+    return th_list[best_axis_index], search_axes[best_axis_index]
 
 
 def qparams_symmetric_selection_histogram_search(error_function: Callable,

{mct-nightly-2.1.0.20240609.524 → mct-nightly-2.1.0.20240611.428}/model_compression_toolkit/core/common/quantization/quantization_params_generation/qparams_weights_computation.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Dict, Any
+from typing import Dict, Any, Tuple
 
 import numpy as np
 
@@ -34,7 +34,7 @@ def get_weights_qparams(weights_attr_values: np.ndarray,
                         output_channels_axis: int,
                         node=None,
                         hessian_info_service: HessianInfoService = None,
-                        num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Dict[Any, Any]:
+                        num_hessian_samples: int = NUM_QPARAM_HESSIAN_SAMPLES) -> Tuple[Dict[Any, Any], int]:
     """
     Compute thresholds to quantize a kernel according to a NodeWeightsQuantizationConfig
     instance.
@@ -50,22 +50,24 @@ def get_weights_qparams(weights_attr_values: np.ndarray,
 
     Returns:
         A dictionary with the quantization threshold of the kernel.
+        Selected quantization channel axis.
     """
     if attr_quant_config.weights_quantization_params_fn is not None:
-        weights_params = attr_quant_config.weights_quantization_params_fn(weights_attr_values,
-                                                                          p=attr_quant_config.l_p_value,
-                                                                          n_bits=attr_quant_config.weights_n_bits,
-                                                                          per_channel=attr_quant_config.weights_per_channel_threshold and output_channels_axis is not None,
-                                                                          channel_axis=output_channels_axis,
-                                                                          min_threshold=weights_quant_config.min_threshold,
-                                                                          quant_error_method=attr_quant_config.weights_error_method,
-                                                                          node=node,
-                                                                          hessian_info_service=hessian_info_service,
-                                                                          num_hessian_samples=num_hessian_samples)
+        weights_params, output_channels_axis = attr_quant_config.weights_quantization_params_fn(
+            weights_attr_values,
+            p=attr_quant_config.l_p_value,
+            n_bits=attr_quant_config.weights_n_bits,
+            per_channel=attr_quant_config.weights_per_channel_threshold,
+            channel_axis=output_channels_axis,
+            min_threshold=weights_quant_config.min_threshold,
+            quant_error_method=attr_quant_config.weights_error_method,
+            node=node,
+            hessian_info_service=hessian_info_service,
+            num_hessian_samples=num_hessian_samples)
     else:
         weights_params = {}
 
-    return weights_params
+    return weights_params, output_channels_axis
 
 
 def _get_kernel_channels_mapping(fw_info:FrameworkInfo,