mct-nightly 2.2.0.20250106.546__tar.gz → 2.2.0.20250107.134735__tar.gz

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.2.0.20250106.546
+Version: 2.2.0.20250107.134735
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.2.0.20250106.546
+Version: 2.2.0.20250107.134735
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/mct_nightly.egg-info/SOURCES.txt

@@ -257,6 +257,7 @@ model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functio
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_layer_norm.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/functional_linear.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/linear_collapsing.py
+model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/matmul_decomposition.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/multi_head_attention_decomposition.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/relu_bound_to_power_of_2.py
 model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/remove_identity.py
@@ -439,6 +440,7 @@ model_compression_toolkit/qat/pytorch/quantizer/ste_rounding/uniform_ste.py
 model_compression_toolkit/target_platform_capabilities/__init__.py
 model_compression_toolkit/target_platform_capabilities/constants.py
 model_compression_toolkit/target_platform_capabilities/immutable.py
+model_compression_toolkit/target_platform_capabilities/tpc_io_handler.py
 model_compression_toolkit/target_platform_capabilities/schema/__init__.py
 model_compression_toolkit/target_platform_capabilities/schema/mct_current_schema.py
 model_compression_toolkit/target_platform_capabilities/schema/schema_functions.py
@@ -457,54 +459,17 @@ model_compression_toolkit/target_platform_capabilities/target_platform/targetpla
 model_compression_toolkit/target_platform_capabilities/tpc_models/__init__.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/get_target_platform_capabilities.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/__init__.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/target_platform_capabilities.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/latest/__init__.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/__init__.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tp_model.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tpc_keras.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1/tpc_pytorch.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/__init__.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/tp_model.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/tpc_keras.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_lut/tpc_pytorch.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/__init__.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tp_model.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_keras.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v1_pot/tpc_pytorch.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/__init__.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tp_model.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_keras.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tpc_pytorch.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/__init__.py
-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/imx500_tpc/v4/__init__.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tp_model.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/tpc_keras.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v4/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
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/v1/__init__.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/v1/tp_model.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/v1/tpc_keras.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/qnnpack_tpc/v1/tpc_pytorch.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/__init__.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/target_platform_capabilities.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/latest/__init__.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/__init__.py
 model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tp_model.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_keras.py
-model_compression_toolkit/target_platform_capabilities/tpc_models/tflite_tpc/v1/tpc_pytorch.py
 model_compression_toolkit/trainable_infrastructure/__init__.py
 model_compression_toolkit/trainable_infrastructure/common/__init__.py
 model_compression_toolkit/trainable_infrastructure/common/annealing_schedulers.py

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/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.2.0.20250106.000546"
+__version__ = "2.2.0.20250107.134735"

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/__init__.py

@@ -18,7 +18,7 @@ from model_compression_toolkit.core.common import network_editors as network_edi
 from model_compression_toolkit.core.common.quantization.debug_config import DebugConfig
 from model_compression_toolkit.core.common.quantization import quantization_config
 from model_compression_toolkit.core.common.mixed_precision import mixed_precision_quantization_config
-from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig, QuantizationErrorMethod, DEFAULTCONFIG
+from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig, QuantizationErrorMethod, DEFAULTCONFIG, CustomOpsetLayers
 from model_compression_toolkit.core.common.quantization.bit_width_config import BitWidthConfig
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py

@@ -49,7 +49,7 @@ def compute_graph_max_cut(memory_graph: MemoryGraph,
     it = 0
     while it < n_iter:
         estimate = (u_bound + l_bound) / 2
-        schedule, max_cut_size, cuts = max_cut_astar.solve(estimate_factor=estimate, iter_limit=astar_n_iter)
+        schedule, max_cut_size, cuts = max_cut_astar.solve(estimate=estimate, iter_limit=astar_n_iter)
         if schedule is None:
             l_bound = estimate
         else:

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/common/graph/memory_graph/cut.py

@@ -64,7 +64,10 @@ class Cut:
         """
         if isinstance(other, Cut):
             return self.mem_elements == other.mem_elements
-        return False
+        return False  # pragma: no cover
 
     def __hash__(self):
-        return hash((frozenset(self.op_order), frozenset(self.op_record), self.mem_elements))
+        return hash((frozenset(self.op_order), frozenset(self.op_record), self.mem_elements))
+
+    def __repr__(self):
+        return f"<Cut: Nodes={[e.node_name for e in self.mem_elements.elements]}, size={self.memory_size()}>"  # pragma: no cover

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/common/graph/memory_graph/max_cut_astar.py

@@ -100,7 +100,7 @@ class MaxCutAstar:
         edges_src_ab = [(src_dummy_a, src_dummy_b)]
         edges_src_ba = [(src_dummy_b, src_a) for src_a in memory_graph.sources_a]
 
-        # Target Cut
+        # Target Cut (Adding 2 consecutive dummy nodes so the final cut will include only dummy tensors).
         target_dummy_a = next(gen_a)
         target_dummy_a2 = next(gen_a)
         target_dummy_b = next(gen_b)
@@ -122,13 +122,13 @@ class MaxCutAstar:
         self.target_cut = Cut([], set(), MemoryElements(elements={target_dummy_b, target_dummy_b2},
                                                         total_size=0))
 
-    def solve(self, estimate_factor: float, iter_limit: int = 500) -> Tuple[List[BaseNode], float, List[Cut]]:
+    def solve(self, estimate: float, iter_limit: int = 500) -> Tuple[List[BaseNode], float, List[Cut]]:
         """
         The AStar solver function. This method runs an AStar-like search on the memory graph,
-        using the given estimate_factor as a heuristic gap for solutions to consider.
+        using the given estimate as a heuristic gap for solutions to consider.
 
         Args:
-            estimate_factor: A multiplication factor which allows the search to consider larger size of nodes in each
+            estimate: Cut size estimation to consider larger size of nodes in each
                 expansion step, in order to fasten the algorithm divergence towards a solution.
             iter_limit: An upper limit for the number of expansion steps that the algorithm preforms.
 
@@ -148,17 +148,14 @@ class MaxCutAstar:
 
         while expansion_count < iter_limit and len(open_list) > 0:
             # Choose next node to expand
-            next_cut = self._get_cut_to_expand(open_list, costs, routes, estimate_factor)
+            next_cut = self._get_cut_to_expand(open_list, costs, routes, estimate)
 
             cut_cost = costs[next_cut]
             cut_route = routes[next_cut]
 
             if next_cut == self.target_cut:
-                # TODO maxcut: Why do we filter the cuts (cut_route) but not the max cut size (cut_sost).
-                #              This is a mismatch between max_cut and max(cuts).
-                #              Also, unfiltered cut_route seems perfect, including input and output tensor sizes of current op.
-                return self._remove_dummys_from_path(cut_route[0].op_order), cut_cost,\
-                       list(set([self._remove_dummys_from_cut(self.clean_memory_for_next_step(c)) for c in cut_route]))
+                return self._remove_dummy_nodes_from_path(cut_route[0].op_order), cut_cost,\
+                       list(set([self._remove_dummy_tensors_from_cut(c) for c in cut_route]))
 
             if self.is_pivot(next_cut):
                 # Can clear all search history
@@ -176,7 +173,7 @@ class MaxCutAstar:
             expansion_count += 1
 
             # Only consider nodes that where not already visited
-            expanded_cuts = list(filter(lambda _c: _c not in closed_list, expanded_cuts))
+            expanded_cuts = [_c for _c in expanded_cuts if _c not in closed_list]
             for c in expanded_cuts:
                 cost = self.accumulate(cut_cost, c.memory_size())
                 if c not in open_list:
@@ -191,7 +188,7 @@ class MaxCutAstar:
                     self._update_expanded_node(c, cost, cut_route, open_list, costs, routes)
 
         # Halt or No Solution
-        # TODO maxcut: this isn't covered in the coverage test. check if needed and remove no cover
+        # TODO maxcut: this isn't covered in the coverage test. Add test and remove no cover
         return None, 0, None  # pragma: no cover
 
     @staticmethod
@@ -214,7 +211,7 @@ class MaxCutAstar:
         routes.update({cut: [cut] + route})
 
     def _get_cut_to_expand(self, open_list: List[Cut], costs: Dict[Cut, float], routes: Dict[Cut, List[Cut]],
-                           estimate_factor: float) -> Cut:
+                           estimate: float) -> Cut:
         """
         An auxiliary method for finding a cut for expanding the search out of a set of potential cuts for expansion.
 
@@ -222,13 +219,15 @@ class MaxCutAstar:
             open_list: The search open list.
             costs: The search utility mapping between cuts and their cost.
             routes: The search utility mapping between cuts and their routes.
-            estimate_factor: A multiplication factor to set extended boundaries on the potential cuts to exapand.
+            estimate: Cut size estimation to set extended boundaries on the potential cuts to expand.
 
         Returns: A sorted list of potential cuts for expansion (ordered by lowest cost first).
 
         """
+        max_cut_len = max([len(routes[c]) for c in open_list])
         ordered_cuts_list = sorted(open_list,
-                                   key=lambda c: (self.accumulate(costs[c], self.estimate(c, estimate_factor)), -len(routes[c])))
+                                   key=lambda c: (self.accumulate(costs[c], self.estimate(c, estimate)),
+                                                  max_cut_len - len(routes[c])))
 
         assert len(ordered_cuts_list) > 0
         return ordered_cuts_list[0]
@@ -356,7 +355,8 @@ class MaxCutAstar:
         # TODO maxcut: this isn't covered in the coverage test. check if needed and remove no cover
         return cost_1 < cost_2  # pragma: no cover
 
-    def estimate(self, cut: Cut, estimate_factor: float) -> float:
+    @staticmethod
+    def estimate(cut: Cut, estimate: float) -> float:
         """
         A function that defines the estimation gap for the Astar search.
         The estimation gap is used to sort the cuts that are considered for expanding the search in each iteration.
@@ -364,15 +364,15 @@ class MaxCutAstar:
         Args:
             cut: A cut (not used in the default implementation, but can be used if overriding the method to consider
                 the actual cut in the estimation computation).
-            estimate_factor: The given estimate factor to the search.
+            estimate: The given estimate to the search.
 
         Returns: An estimation value.
 
         """
-        return estimate_factor * self.memory_graph.memory_lbound_single_op
+        return estimate
 
     @staticmethod
-    def get_init_estimate_factor(memory_graph: MemoryGraph) -> float:
+    def get_init_estimate(memory_graph: MemoryGraph) -> float:  # pragma: no cover
         """
         Returns an initial estimation value, which is based on the memory graph's upper and lower bounds.
 
@@ -383,12 +383,12 @@ class MaxCutAstar:
 
         """
         # TODO maxcut: this isn't covered in the coverage test. check if needed and remove no cover
-        l_bound = memory_graph.memory_lbound_single_op  # pragma: no cover
-        u_bound = 2 * sum([t.total_size for t in memory_graph.b_nodes]) - l_bound  # pragma: no cover
-        return (u_bound + l_bound) / 2  # pragma: no cover
+        l_bound = memory_graph.memory_lbound_single_op
+        u_bound = 2 * sum([t.total_size for t in memory_graph.b_nodes]) - l_bound
+        return (u_bound + l_bound) / 2
 
     @staticmethod
-    def _remove_dummys_from_path(path: List[BaseNode]) -> List[BaseNode]:
+    def _remove_dummy_nodes_from_path(path: List[BaseNode]) -> List[BaseNode]:
         """
         An auxiliary method which removes dummy nodes from a given list of nodes (a path in the graph).
 
@@ -401,7 +401,7 @@ class MaxCutAstar:
         return list(filter(lambda n: DUMMY_NODE not in n.name, path))
 
     @staticmethod
-    def _remove_dummys_from_cut(cut: Cut) -> Cut:
+    def _remove_dummy_tensors_from_cut(cut: Cut) -> Cut:
         """
         An auxiliary method which removes dummy nodes from a given cut.
 
@@ -411,7 +411,7 @@ class MaxCutAstar:
         Returns: The same cut without dummy nodes elements.
 
         """
-        filtered_memory_elements = set(filter(lambda elm: DUMMY_TENSOR not in elm.node_name, cut.mem_elements.elements))
+        filtered_memory_elements = set([elm for elm in cut.mem_elements.elements if DUMMY_TENSOR not in elm.node_name])
 
         return Cut(cut.op_order, cut.op_record,
                    mem_elements=MemoryElements(filtered_memory_elements,

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/common/quantization/quantization_config.py

@@ -13,13 +13,30 @@
 # limitations under the License.
 # ==============================================================================
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 import math
 from enum import Enum
+from typing import Optional, Dict, NamedTuple, List
 
+from model_compression_toolkit import DefaultDict
 from model_compression_toolkit.constants import MIN_THRESHOLD
 
 
+class CustomOpsetLayers(NamedTuple):
+    """
+    This struct defines a set of operators from a specific framework, which will be used to configure a custom operator
+    set in the TPC.
+
+    Args:
+        operators: a list of framework operators to map to a certain custom opset name.
+        attr_mapping: a mapping between an opset name to a mapping between its attributes' general names and names in
+            the framework.
+    """
+
+    operators: List
+    attr_mapping: Optional[Dict[str, DefaultDict]] = None
+
+
 class QuantizationErrorMethod(Enum):
     """
     Method for quantization threshold selection:
@@ -86,6 +103,7 @@ class QuantizationConfig:
     concat_threshold_update: bool = False
     activation_bias_correction: bool = False
     activation_bias_correction_threshold: float = 0.0
+    custom_tpc_opset_to_layer: Optional[Dict[str, CustomOpsetLayers]] = None
 
 
 # Default quantization configuration the library use.

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/keras/back2framework/keras_model_builder.py

@@ -49,38 +49,6 @@ FQ_NODE_OP_V2_4 = 'quantization.fake_quant_with_min_max_vars'
 BATCH_INPUT_SHAPE = 'batch_input_shape'
 
 
-def get_node_name_from_layer(layer: Layer) -> str:
-    """
-    Get a node's name from the layer it was built from. For TensorFlowOpLayer
-    we remove the prefix "tf_op_layer".
-
-    Args:
-        layer: Keras Layer to get its corresponding node's name.
-
-    Returns:
-        Name of the node that was built from the passed layer.
-    """
-
-    name = layer.name
-    if isinstance(layer, TensorFlowOpLayer):  # remove TF op layer prefix
-        name = '_'.join(name.split('_')[3:])
-    return name
-
-
-def is_layer_fake_quant(layer: Layer) -> bool:
-    """
-    Check whether a Keras layer is a fake quantization layer or not.
-    Args:
-        layer: Layer to check if it's a fake quantization layer or not.
-
-    Returns:
-        Whether a Keras layer is a fake quantization layer or not.
-    """
-    # in tf2.3 fake quant node is implemented as TensorFlowOpLayer, while in tf2.4 as TFOpLambda
-    return (isinstance(layer, TensorFlowOpLayer) and layer.node_def.op == FQ_NODE_OP_V2_3) or (
-            isinstance(layer, TFOpLambda) and layer.symbol == FQ_NODE_OP_V2_4)
-
-
 class KerasModelBuilder(BaseModelBuilder):
     """
     Builder for Keras models.
@@ -291,7 +259,7 @@ class KerasModelBuilder(BaseModelBuilder):
                                     arg = n.weights.get(pos)
                                     if arg is None:
                                         if len(input_tensors) == 0:
-                                            Logger.critical(f"Couldn't find a weight or input tensor matching operator's "
+                                            Logger.critical(f"Couldn't find a weight or input tensor matching operator's "  # pragma: no cover
                                                             f"argument name '{k}' in location {pos} for node {n.name}.")
                                         arg = input_tensors.pop(0)
                                     op_call_kwargs.update({k: arg})

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/keras/graph_substitutions/substitutions/conv_funcs_to_layer.py

@@ -141,7 +141,7 @@ class Conv2dFuncToConv2dLayer(common.BaseSubstitution):
         strides = self._parse_tf_stride_dilation(conv_func_node, STRIDES)
         if strides is None:
             # Non-standard strides -> skip substitution.
-            return graph
+            return graph  # pragma: no cover
         conv_fw_attr[STRIDES] = strides
 
         padding = conv_func_node.op_call_kwargs.get(PADDING) or 'VALID'
@@ -153,7 +153,7 @@ class Conv2dFuncToConv2dLayer(common.BaseSubstitution):
         dilations = self._parse_tf_stride_dilation(conv_func_node, DILATIONS)
         if dilations is None:
             # Non-standard dilations -> skip substitution.
-            return graph
+            return graph  # pragma: no cover
         conv_fw_attr[DILATION_RATE] = dilations
 
         if b is None:

{mct_nightly-2.2.0.20250106.546 → mct_nightly-2.2.0.20250107.134735}/model_compression_toolkit/core/keras/resource_utilization_data_facade.py

@@ -18,6 +18,7 @@ from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, Cor
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
+from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformModel
 from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_data import compute_resource_utilization_data
 from model_compression_toolkit.verify_packages import FOUND_TF
@@ -27,6 +28,8 @@ if FOUND_TF:
     from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
     from tensorflow.keras.models import Model
+    from model_compression_toolkit.target_platform_capabilities.target_platform.targetplatform2framework.attach2keras import \
+        AttachTpcToKeras
 
     from model_compression_toolkit import get_target_platform_capabilities
 
@@ -36,7 +39,8 @@ if FOUND_TF:
                                         representative_data_gen: Callable,
                                         core_config: CoreConfig = CoreConfig(
                                             mixed_precision_config=MixedPrecisionQuantizationConfig()),
-                                        target_platform_capabilities: TargetPlatformCapabilities = KERAS_DEFAULT_TPC) -> ResourceUtilization:
+                                        target_platform_capabilities: TargetPlatformModel = KERAS_DEFAULT_TPC
+                                        ) -> ResourceUtilization:
         """
         Computes resource utilization data that can be used to calculate the desired target resource utilization
         for mixed-precision quantization.
@@ -78,6 +82,12 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        # Attach tpc model to framework
+        attach2keras = AttachTpcToKeras()
+        target_platform_capabilities = attach2keras.attach(
+            target_platform_capabilities,
+            custom_opset2layer=core_config.quantization_config.custom_tpc_opset_to_layer)
+
         return compute_resource_utilization_data(in_model,
                                                  representative_data_gen,
                                                  core_config,