mct-nightly 1.11.0.20240317.91316__py3-none-any.whl → 1.11.0.20240319.407__py3-none-any.whl

{mct_nightly-1.11.0.20240317.91316.dist-info → mct_nightly-1.11.0.20240319.407.dist-info}/METADATA

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

{mct_nightly-1.11.0.20240317.91316.dist-info → mct_nightly-1.11.0.20240319.407.dist-info}/RECORD

@@ -1,13 +1,13 @@
-model_compression_toolkit/__init__.py,sha256=iJs169mTUA7Yi3p0vgbTNRoyzRYnpn7W-AhLFWeJwUE,1574
+model_compression_toolkit/__init__.py,sha256=xX48bBE90QORHpae8fqX4wNK9NTEi5qT_JQoOG7-tPc,1574
 model_compression_toolkit/constants.py,sha256=_OW_bUeQmf08Bb4oVZ0KfUt-rcCeNOmdBv3aP7NF5fM,3631
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=b9DVktZ-LymFcRxv2aL_sdiE6S2sSrFGWltx6dgEuUY,4863
-model_compression_toolkit/core/__init__.py,sha256=P-7OYR4TFYxVV_ZpIJBogkX8bGvXcijlF65Ez3ivjhc,1838
+model_compression_toolkit/core/__init__.py,sha256=DRw7VF7jsqHxGtoxf8F0YcXPRRnQIw6sn6Q925MmWC8,1944
 model_compression_toolkit/core/analyzer.py,sha256=dbsD61pakp_9JXNyAScLdtJvcXny9jr_cMbET0Bd3Sg,2975
 model_compression_toolkit/core/exporter.py,sha256=Zo_C5GjIzihtJOyGp-xeCVhY_qohkVz_EGyrSZCbWRM,4115
-model_compression_toolkit/core/graph_prep_runner.py,sha256=3xp0WYqyeRdlBkf5R6uD2zWubg_JPttOwS7JRhKykBY,10043
+model_compression_toolkit/core/graph_prep_runner.py,sha256=Ftqm59hT5TGWmSNkY9bFZkVfCacpGyZfCe-6yZR5WY0,10100
 model_compression_toolkit/core/quantization_prep_runner.py,sha256=hFhDkS8GwzXZ7Ho_9qbbb8DAAWs3OONOfMSD5OU_b0o,6153
-model_compression_toolkit/core/runner.py,sha256=hXnbgP8Q-62Ie4wAq4JXO-2o77uR3le4mHYgFqJOvfc,10928
+model_compression_toolkit/core/runner.py,sha256=FJ_TG-OZDtDBM_BNHTdcAX5NKAWPog-0Gh3uDgnXUxU,11383
 model_compression_toolkit/core/common/__init__.py,sha256=Wh127PbXcETZX_d1PQqZ71ETK3J9XO5A-HpadGUbj6o,1447
 model_compression_toolkit/core/common/base_substitutions.py,sha256=xDFSmVVs_iFSZfajytI0cuQaNRNcwHX3uqOoHgVUvxQ,1666
 model_compression_toolkit/core/common/data_loader.py,sha256=jCoVIb4yeOWyCrCNRB1W-mgLSyqNVGEepFXrIqufVc4,4119
@@ -62,8 +62,8 @@ model_compression_toolkit/core/common/mixed_precision/bit_width_setter.py,sha256
 model_compression_toolkit/core/common/mixed_precision/configurable_quant_id.py,sha256=LLDguK7afsbN742ucLpmJr5TUfTyFpK1vbf2bpVr1v0,882
 model_compression_toolkit/core/common/mixed_precision/configurable_quantizer_utils.py,sha256=7dKMi5S0zQZ16m8NWn1XIuoXsKuZUg64G4-uK8-j1PQ,5177
 model_compression_toolkit/core/common/mixed_precision/distance_weighting.py,sha256=H8qYkJsk88OszUJo-Zde7vTmWiypLTg9KbbzIZ-hhvM,2812
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=4skskrRuKOoMZIX9XB9Os3WmQiFq8rEe05RmL6xjrxo,4553
-model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=sXyhmyO0FtGO7vHOWqUign88Kh7MCzN8Ohk6wIXq0GQ,6992
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py,sha256=DP5tcxPtiVbSWAeoFbEp7iTwpxDBU1g7V5w7ehDG6jI,4573
+model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py,sha256=mN-QeabIu_Mz1IzPeQjqgqprCTdwGm4ThYX0gZAek-E,7103
 model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_manager.py,sha256=f5j7R7A_bSVqeBY4WuDN8n0YWlR8jhK_n9eKInQ8anY,36763
 model_compression_toolkit/core/common/mixed_precision/sensitivity_evaluation.py,sha256=vt829yxXlfbQHPDUHJebda7jfzpGf1N3b6L4XJ4zbSI,28534
 model_compression_toolkit/core/common/mixed_precision/set_layer_to_bitwidth.py,sha256=P8QtKgFXtt5b2RoubzI5OGlCfbEfZsAirjyrkFzK26A,2846
@@ -71,7 +71,7 @@ model_compression_toolkit/core/common/mixed_precision/solution_refinement_proced
 model_compression_toolkit/core/common/mixed_precision/kpi_tools/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi.py,sha256=gcwwuzLKpa2WvsyAr6MXb4cXhOxCM0dvVHKLL-FWGoA,4297
 model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi_aggregation_methods.py,sha256=X0PbF3UHVy3JRRIgcogKpTNm26AJOJ7blajAWsDf7R4,3920
-model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi_data.py,sha256=vrgUYeL6MPVO_tBNIGf6tuOXsjl60JcIW8y_CPPk6Tk,7464
+model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi_data.py,sha256=7wzhz7tPMVq-JmcHvTGiOdDQPXNw3i7HXAvbymJOXY8,7618
 model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi_functions_mapping.py,sha256=cjLf_g4n1INlT1TE1z-I41hDXUTTy8krUSvhRB57rv0,1602
 model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi_methods.py,sha256=yg8Pg9kMsjQzi03tcoQlp0iqnnqKvxjohOQNzPdRPzs,20840
 model_compression_toolkit/core/common/mixed_precision/search_methods/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
@@ -86,9 +86,9 @@ model_compression_toolkit/core/common/pruning/greedy_mask_calculator.py,sha256=G
 model_compression_toolkit/core/common/pruning/memory_calculator.py,sha256=RnMmgNDHekKFOj-b-ad5rhjuKUvbVawy1A31nxuCRTg,19217
 model_compression_toolkit/core/common/pruning/prune_graph.py,sha256=ddbZLuWvlNoj5so_5NRbIuG5qDFxD9ApG2gPirbov8o,3317
 model_compression_toolkit/core/common/pruning/pruner.py,sha256=vXxzBXQ-oAEnw6PAD1SUiNXX7Xix4JJ0LAmV04sjFz0,7313
-model_compression_toolkit/core/common/pruning/pruning_config.py,sha256=IfF824hNttyw2i4Tuf3g8CUfelJR3eZuOLzf2aEZNAM,3442
+model_compression_toolkit/core/common/pruning/pruning_config.py,sha256=PO4C1C1_hhAX_B05kqpC-TTx1S1O6Dj9DrtZrxpi1aE,3670
 model_compression_toolkit/core/common/pruning/pruning_framework_implementation.py,sha256=H2gnCv-lyRLXapDy71QHA3JkLYTQT1ni23nGTYErsZo,6734
-model_compression_toolkit/core/common/pruning/pruning_info.py,sha256=DK_ofX-73tfpdmkHNLYlO6_SBifZDsRWmGHsCbrUFN8,4083
+model_compression_toolkit/core/common/pruning/pruning_info.py,sha256=gSCh_qXmLATChb5Nh16wvR1ffI9SERstEroc_hFrVQo,3781
 model_compression_toolkit/core/common/pruning/pruning_section.py,sha256=I4vxh5iPKWs8yji-q4TVmaa6bcfLwT3ZhB2x8F8cJkU,5721
 model_compression_toolkit/core/common/pruning/importance_metrics/__init__.py,sha256=3Lkr37Exk9u8811hw8hVqkGcbTQGcLjd3LLuLC3fa_E,698
 model_compression_toolkit/core/common/pruning/importance_metrics/base_importance_metric.py,sha256=qMAtLWs5fjbSco8nhbig5TkuacdhnDW7cy3avMHRGX4,1988
@@ -99,7 +99,7 @@ model_compression_toolkit/core/common/pruning/mask/per_channel_mask.py,sha256=hk
 model_compression_toolkit/core/common/pruning/mask/per_simd_group_mask.py,sha256=gmzD32xsfJH8vkkqaspS7vYa6VWayk1GJe-NfoAEugQ,5901
 model_compression_toolkit/core/common/quantization/__init__.py,sha256=sw7LOPN1bM82o3SkMaklyH0jw-TLGK0-fl2Wq73rffI,697
 model_compression_toolkit/core/common/quantization/candidate_node_quantization_config.py,sha256=D7lgCc0drQZ3yXNctTBg-FnqHX7e32zp0-ocGYGJbEE,4553
-model_compression_toolkit/core/common/quantization/core_config.py,sha256=IkD4Jl9PWdPucfUMq0TtyUl5DBJvha7Dd2xSW7_7dz8,2015
+model_compression_toolkit/core/common/quantization/core_config.py,sha256=KYdyfSmjSL4ye24nKlC_c4_AxYb14qoqaeMnZj4-8kE,2257
 model_compression_toolkit/core/common/quantization/debug_config.py,sha256=HtkMmneN-EmAzgZK4Vp4M8Sqm5QKdrvNyyZMpaVqYzY,1482
 model_compression_toolkit/core/common/quantization/filter_nodes_candidates.py,sha256=fwF4VILaX-u3ZaFd81xjbJuhg8Ef-JX_KfMXW0TPV-I,7136
 model_compression_toolkit/core/common/quantization/node_quantization_config.py,sha256=HWBBF--cbzsiMx3BG2kQ3JHkfalVnGO3N-rAXMwNqp4,26707
@@ -108,7 +108,7 @@ model_compression_toolkit/core/common/quantization/quantization_fn_selection.py,
 model_compression_toolkit/core/common/quantization/quantization_params_fn_selection.py,sha256=mrgVzZszWjxnjT8zm77UVLWKTOwd2thGBo6WNqAS4X8,3867
 model_compression_toolkit/core/common/quantization/quantize_graph_weights.py,sha256=xnM9O9LshYw3dprqfsnK9mw7ipOEAkI85o20auyfswg,2626
 model_compression_toolkit/core/common/quantization/quantize_node.py,sha256=cdzGNWfT4MRogIU8ehs0tr3lVjnzAI-jeoS9b4TwVBo,2854
-model_compression_toolkit/core/common/quantization/set_node_quantization_config.py,sha256=onzV581FPw19WN00EUjrXkvPbA6msHk-VSJkBVduV-s,11490
+model_compression_toolkit/core/common/quantization/set_node_quantization_config.py,sha256=ntfdEK39SAuegHtGa_v9H-_IC9WagRvwokRL3wEnGso,11491
 model_compression_toolkit/core/common/quantization/quantization_params_generation/__init__.py,sha256=eCDGwsWYLU6z7qbEVb4TozMW_nd5VEP_iCJ6PcvyEPw,1486
 model_compression_toolkit/core/common/quantization/quantization_params_generation/error_functions.py,sha256=eH3nSXPFn94ATF3dZn2HxNAGVJUWotirN6o8wwDfkLg,18165
 model_compression_toolkit/core/common/quantization/quantization_params_generation/lut_kmeans_params.py,sha256=h8Zmpq3KdcsdUUy7K1fvWOVSki0mxT8wtKZXGmgFl74,7405
@@ -211,7 +211,7 @@ model_compression_toolkit/core/keras/visualization/__init__.py,sha256=mjbqLD-KcG
 model_compression_toolkit/core/pytorch/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/pytorch/constants.py,sha256=NI-J7REuxn06oEIHsmJ4GqtNC3TbV8xlkJjt5Ar-c4U,2626
 model_compression_toolkit/core/pytorch/default_framework_info.py,sha256=r1XyzUFvrjGcJHQM5ETLsMZIG2yHCr9HMjqf0ti9inw,4175
-model_compression_toolkit/core/pytorch/kpi_data_facade.py,sha256=fMUUHOv31FGWy1dUXteWtj6OlVm4QC2mf2H77n7ToLM,4584
+model_compression_toolkit/core/pytorch/kpi_data_facade.py,sha256=eKFq0gO2DrlS_wN4plMGZTabQKmb0pylIVGZ44HqSnw,4527
 model_compression_toolkit/core/pytorch/pytorch_device_config.py,sha256=IoMvTch5awAEPvB6Tg6ANhFGXvfSgv7JLsUBlxpMwk4,4330
 model_compression_toolkit/core/pytorch/pytorch_implementation.py,sha256=1uIDT-3wLzQf1FT8fMleyu5w5EYL0n7HoFEG80XDUY8,27082
 model_compression_toolkit/core/pytorch/pytorch_node_prior_info.py,sha256=n_B4a6FMwM9D2w8kzy3oenBWZgXNZuIZgTJC6JEuTy0,3250
@@ -263,19 +263,19 @@ model_compression_toolkit/core/pytorch/reader/node_holders.py,sha256=TaolORuwBZE
 model_compression_toolkit/core/pytorch/reader/reader.py,sha256=Co3-AHZCEOw5w-jtgf9oAKsgtjQoG0MeeSeBVnQ0xOA,5801
 model_compression_toolkit/core/pytorch/statistics_correction/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/pytorch/statistics_correction/apply_second_moment_correction.py,sha256=VgU24J3jf7QComHH7jonOXSkg6mO4TOch3uFkOthZvM,3261
-model_compression_toolkit/data_generation/__init__.py,sha256=zp3nQ7NhDncuGdHBwCXkRJh6JnGoTYhZZlAOrDE8omc,1138
+model_compression_toolkit/data_generation/__init__.py,sha256=R_RnB8Evj4uq0WKiPWvBWfeePrbake7Z03ugJgK7jLo,1466
 model_compression_toolkit/data_generation/common/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/common/constants.py,sha256=21e3ZX9WVYojexG2acTgklrBk8ZO9DjJnKpP4KHZC44,1018
 model_compression_toolkit/data_generation/common/data_generation.py,sha256=PnKkWCBf4yla0E4LhvOqT8htWiGW4F98bygExQnpwqI,6397
 model_compression_toolkit/data_generation/common/data_generation_config.py,sha256=ynyNaT2x2d23bYSrO2sRItM2ZsjGD0K0fM71FlibiJQ,4564
-model_compression_toolkit/data_generation/common/enums.py,sha256=UJhndTsE7q7Bm6CgCYQKWOuuD-9lj6j_QQ28KWZK8uU,3522
+model_compression_toolkit/data_generation/common/enums.py,sha256=OGnvtEGFbP5l4V3-1l32zzVQwTb1vGJhTVF0kOkYZK4,3584
 model_compression_toolkit/data_generation/common/image_pipeline.py,sha256=WwyeoIvgmcxKnuOX-_Hl_0APET4M26f5x-prhUB3qvU,2149
 model_compression_toolkit/data_generation/common/model_info_exctractors.py,sha256=9zYlyuc7K1s2neHWF3wqL5EVOVaoz_QkCYTktSXrSXI,6047
 model_compression_toolkit/data_generation/common/optimization_utils.py,sha256=8wCU-bCLabDIUayG3eyapdD8pTE6x0RYG5o3rfha7XE,19572
 model_compression_toolkit/data_generation/keras/__init__.py,sha256=lNJ29DYxaLUPDstRDA1PGI5r9Fulq_hvrZMlhst1Z5g,697
 model_compression_toolkit/data_generation/keras/constants.py,sha256=uy3eU24ykygIrjIvwOMj3j5euBeN2PwWiEFPOkJJ7ss,1088
 model_compression_toolkit/data_generation/keras/image_pipeline.py,sha256=_Qezq67huKmmNsxdFBBrTY-VaGR-paFzDH80dDuRnug,7623
-model_compression_toolkit/data_generation/keras/keras_data_generation.py,sha256=MYFdMPqGxy9tRaTIstJMkcYOk0tMXirke5fxdIJvBjU,19720
+model_compression_toolkit/data_generation/keras/keras_data_generation.py,sha256=6UQXrpwghxhHgMOGP8he84y4ZUQj2v0UGFRNMIPhBI8,21587
 model_compression_toolkit/data_generation/keras/model_info_exctractors.py,sha256=b3BaOGiMAlCCzPICww722l2H_RucoHgpGUK6xYe8xTA,8552
 model_compression_toolkit/data_generation/keras/optimization_utils.py,sha256=uQAJpJPpnLDTTLDQGyTS0ZYp2T38TTZLOOElcJPBKHA,21146
 model_compression_toolkit/data_generation/keras/optimization_functions/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
@@ -289,7 +289,7 @@ model_compression_toolkit/data_generation/pytorch/constants.py,sha256=QWyreMImcf
 model_compression_toolkit/data_generation/pytorch/image_pipeline.py,sha256=6g7OpOuO3cU4TIuelaRjBKpCPgiMbe1a3iy9bZtdZUo,6617
 model_compression_toolkit/data_generation/pytorch/model_info_exctractors.py,sha256=wxtaQad4aP8D0SgA8qEPORZM3qBD22G6zO1gjwTNIVU,9632
 model_compression_toolkit/data_generation/pytorch/optimization_utils.py,sha256=AjYsO-lm06JOUMoKkS6VbyF4O_l_ffWXrgamqJm1ofE,19085
-model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=BCJ6PVncBBm6sa4IWCYvC-U0-XPs7LV-deao0lq_D20,19192
+model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=Fj0pZEdbQUyydbPRrUWAK9X3hJRtzeFkQ1kCsxHjW84,21012
 model_compression_toolkit/data_generation/pytorch/optimization_functions/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/pytorch/optimization_functions/batchnorm_alignment_functions.py,sha256=dMc4zz9XfYfAT4Cxns57VgvGZWPAMfaGlWLFyCyl8TA,1968
 model_compression_toolkit/data_generation/pytorch/optimization_functions/bn_layer_weighting_functions.py,sha256=i3ePEI8xDE3xZEtmzT5lCkLn9wpObUi_OgqnVDf7nj8,2597
@@ -300,14 +300,14 @@ model_compression_toolkit/exporter/__init__.py,sha256=Eg3c4EAjW3g6h13A-Utgf9ncHr
 model_compression_toolkit/exporter/model_exporter/__init__.py,sha256=9HIBmj8ROdCA-yvkpA8EcN6RHJe_2vEpLLW_gxOJtak,698
 model_compression_toolkit/exporter/model_exporter/fw_agonstic/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/exporter/model_exporter/fw_agonstic/exporter.py,sha256=eSC6gEMc9KY5EwVRam9pJCBpCm0ksUeobKV_JAOap9M,2017
-model_compression_toolkit/exporter/model_exporter/fw_agonstic/quantization_format.py,sha256=oRlw_LG0UNi8Vl77E1U6WDBsEp1eNaphUbG9rSZkSl4,800
+model_compression_toolkit/exporter/model_exporter/fw_agonstic/quantization_format.py,sha256=otuyY3N2h6NmZKjptRvHEnwJRkPVJ2Ty20J1Mwbkjqc,1165
 model_compression_toolkit/exporter/model_exporter/keras/__init__.py,sha256=uZ2RigbY9O2PJ0Il8wPpS_s7frgg9WUGd_SHeKGyl1A,699
 model_compression_toolkit/exporter/model_exporter/keras/base_keras_exporter.py,sha256=-wr2n0yRlmFixXBeZuxg6Rzlvz-ZFUX-PJgSXhgMrEo,1593
 model_compression_toolkit/exporter/model_exporter/keras/export_serialization_format.py,sha256=v_-rOsWDFI-3k8CoJIr-XzT7ny8WXpAMteWRWtTzaeg,963
 model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_keras_exporter.py,sha256=E_1IqFYAGUMOrt3U_JK1k--8D0WzWPbjZH_IRLGw_wY,11478
 model_compression_toolkit/exporter/model_exporter/keras/fakely_quant_tflite_exporter.py,sha256=sqzqQ8US24WgDbg_FoP1NQBgqCbSVwrVTWrxcyY0nPA,3514
 model_compression_toolkit/exporter/model_exporter/keras/int8_tflite_exporter.py,sha256=nGtpDTeH5Tdp7sjyuXsy_9TPpijDYp4nkz366DUUJ0Q,8048
-model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py,sha256=O-GApieS7_zLkpygnN0YvDK-HkCChwA4bSExbI5jvQ8,5998
+model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py,sha256=EByj03xnyJTIwopuKjKne9Nwdr0VqoNdQOCEFwSeTNw,5792
 model_compression_toolkit/exporter/model_exporter/keras/mctq_keras_exporter.py,sha256=qXXkv3X_wb7t622EOHwXIxfGLGaDqh0T0y4UxREi4Bo,1976
 model_compression_toolkit/exporter/model_exporter/pytorch/__init__.py,sha256=uZ2RigbY9O2PJ0Il8wPpS_s7frgg9WUGd_SHeKGyl1A,699
 model_compression_toolkit/exporter/model_exporter/pytorch/base_pytorch_exporter.py,sha256=UPVkEUQCMZ4Lld6CRnEOPEmlfe5vcQZG0Q3FwRBodD4,4021
@@ -331,7 +331,7 @@ model_compression_toolkit/exporter/model_wrapper/pytorch/builder/node_to_quantiz
 model_compression_toolkit/gptq/__init__.py,sha256=YKg-tMj9D4Yd0xW9VRD5EN1J5JrmlRbNEF2fOSgodqA,1228
 model_compression_toolkit/gptq/runner.py,sha256=MIg-oBtR1nbHkexySdCJD_XfjRoHSknLotmGBMuD5qM,5924
 model_compression_toolkit/gptq/common/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/gptq/common/gptq_config.py,sha256=U33sLIPB0pI4h_zhr4X_S9K0cEJWTbWFxkj8z9IGlxg,5268
+model_compression_toolkit/gptq/common/gptq_config.py,sha256=6xP99B-lK1bwGv3AdqxnW1V51z2VdzQcjvoSgJOmygA,5288
 model_compression_toolkit/gptq/common/gptq_constants.py,sha256=QSm6laLkIV0LYmU0BLtmKp3Fi3SqDfbncFQWOGA1cGU,611
 model_compression_toolkit/gptq/common/gptq_framework_implementation.py,sha256=n3mSf4J92kFjekzyGyrJULylI-8Jf5OVWJ5AFoVnEx0,1266
 model_compression_toolkit/gptq/common/gptq_graph.py,sha256=8qmty-2MzV6USRoHgShCA13HqxDI3PDGJaFKCQPFo5E,3026
@@ -341,7 +341,7 @@ model_compression_toolkit/gptq/keras/gptq_keras_implementation.py,sha256=axBwnCS
 model_compression_toolkit/gptq/keras/gptq_loss.py,sha256=rbRkF15MYd6nq4G49kcjb_dPTa-XNq9cTkrb93mXawo,6241
 model_compression_toolkit/gptq/keras/gptq_training.py,sha256=cASZlTmnth3Vu-7GfmC03FxWSXtpSVhdPKT_twWml68,17949
 model_compression_toolkit/gptq/keras/graph_info.py,sha256=FIGqzJbG6GkdHenvdMu-tGTjp4j9BewdF_spmWCb4Mo,4627
-model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=wRyQrJJ71JwtFoiIdBPDHE0srpUwmL7nqHbXOvjDHFc,13578
+model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=NyVMvDsgxMsaAtYxYwaqeQX3VD5GmfftXKHV5wUBLZg,13762
 model_compression_toolkit/gptq/keras/quantizer/__init__.py,sha256=-DK1CDXvlsnEbki4lukZLpl6Xrbo91_jcqxXlG5Eg6Q,963
 model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py,sha256=8NrJBftKFbMAF_jYaAbLP6GBwpCv3Ln1NKURaV75zko,4770
 model_compression_toolkit/gptq/keras/quantizer/quant_utils.py,sha256=Vt7Qb8i4JsE4sFtcjpfM4FTXTtfV1t6SwfoNH8a_Iaw,5055
@@ -358,7 +358,7 @@ model_compression_toolkit/gptq/pytorch/gptq_loss.py,sha256=kDuWw-6zh17wZpYWh4Xa9
 model_compression_toolkit/gptq/pytorch/gptq_pytorch_implementation.py,sha256=tECPTavxn8EEwgLaP2zvxdJH6Vg9jC0YOIMJ7857Sdc,1268
 model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=9zQC42RfAj4ak-XOzF8xEXS3IkHKhKlOClIfaUA0bGI,15396
 model_compression_toolkit/gptq/pytorch/graph_info.py,sha256=-0GDC2cr-XXS7cTFTnDflJivGN7VaPnzVPsxCE-vZNU,3955
-model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=MMp97lTatr0moe0r4cycqNw-1qVo_ixvissH6n_wjnE,12091
+model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=ER5VPSkZZjqYj7PJ-3B5RX33YjHz3tJ4Er9SF6M-93c,12369
 model_compression_toolkit/gptq/pytorch/quantizer/__init__.py,sha256=ZHNHo1yzye44m9_ht4UUZfTpK01RiVR3Tr74-vtnOGI,968
 model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py,sha256=Zb-P0yRyZHHBlDvUBdRwxDpdduEJyJp6OT9pfKFF5ks,4171
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
@@ -372,20 +372,20 @@ model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/__init__.py,sha256
 model_compression_toolkit/gptq/pytorch/quantizer/ste_rounding/symmetric_ste.py,sha256=6uxq_w62jn8DDOt9T7VtA6jZ8jTAPcbTufKFOYpVUm4,8768
 model_compression_toolkit/pruning/__init__.py,sha256=lQMZS8G0pvR1LVi53nnJHNXgLNTan_MWMdwsVxhjrow,1106
 model_compression_toolkit/pruning/keras/__init__.py,sha256=3Lkr37Exk9u8811hw8hVqkGcbTQGcLjd3LLuLC3fa_E,698
-model_compression_toolkit/pruning/keras/pruning_facade.py,sha256=B2mkCh3_AKc1O3IBOdo03PuIyjAoK3IBmgBdmIfUkDI,8296
+model_compression_toolkit/pruning/keras/pruning_facade.py,sha256=2vMmI9QaH9nReyDqZKiWOZPQC3HUQ2ZCahHIMFyveMQ,8396
 model_compression_toolkit/pruning/pytorch/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
 model_compression_toolkit/pruning/pytorch/pruning_facade.py,sha256=ZLmMhwAEnbXNRwMwgoGEGNmHpZx_KWYu7yi5K3aICWI,9184
 model_compression_toolkit/ptq/__init__.py,sha256=Z_hkmTh7aLFei1DJKV0oNVUbrv_Q_0CTw-qD85Xf8UM,904
 model_compression_toolkit/ptq/runner.py,sha256=_c1dSjlPPpsx59Vbg1buhG9bZq__OORz1VlPkwjJzoc,2552
 model_compression_toolkit/ptq/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=6kZ54SY_Slw2DGzALm7X2TzZEej9-FEoaVkjxOdFxp8,8598
+model_compression_toolkit/ptq/keras/quantization_facade.py,sha256=ergUI8RDA2h4_SHU05x2pYJatt-U-fZUrShdHJDLo_o,8844
 model_compression_toolkit/ptq/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=5QH-khLMPFLdCUPQCxOCYY5v4p_M67TZcfCZGnsWqVs,7191
+model_compression_toolkit/ptq/pytorch/quantization_facade.py,sha256=WKzokgg_gGcEHipVH26shneiAiTdSa7d_UUQKoS8ALY,7438
 model_compression_toolkit/qat/__init__.py,sha256=kj2qsZh_Ca7PncsHKcaL5EVT2H8g4hYtvaQ3KFxOkwE,1143
 model_compression_toolkit/qat/common/__init__.py,sha256=6tLZ4R4pYP6QVztLVQC_jik2nES3l4uhML0qUxZrezk,829
 model_compression_toolkit/qat/common/qat_config.py,sha256=zoq0Vb74vCY7WlWD8JH_KPrHDoUHSvMc3gcO53u7L2U,3394
 model_compression_toolkit/qat/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/keras/quantization_facade.py,sha256=xH05Ro9aY9HabQo_PztaXw0-D3Cxvl-GYCmDKRjwkuI,16524
+model_compression_toolkit/qat/keras/quantization_facade.py,sha256=9qWdNIIx2hKmjGCpSGGEAv7HXg91Y9ZuyKE-avHn46c,16784
 model_compression_toolkit/qat/keras/quantizer/__init__.py,sha256=zmYyCa25_KLCSUCGUDRslh3RCIjcRMxc_oXa54Aui-4,996
 model_compression_toolkit/qat/keras/quantizer/base_keras_qat_quantizer.py,sha256=gPuIgQb8OafvC3SuA8jNsGoy8S8eTsDCEKuh36WDNss,2104
 model_compression_toolkit/qat/keras/quantizer/quant_utils.py,sha256=cBULOgWUodcBO1lHevZggdTevuDYI6tQceV86U2x6DA,2543
@@ -397,7 +397,7 @@ model_compression_toolkit/qat/keras/quantizer/ste_rounding/__init__.py,sha256=cc
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/symmetric_ste.py,sha256=I4KlaGv17k71IyjuSG9M0OlXlD5P0pfvKa6oCyRQ5FE,13517
 model_compression_toolkit/qat/keras/quantizer/ste_rounding/uniform_ste.py,sha256=EED6LfqhX_OhDRJ9e4GwbpgNC9vq7hoXyJS2VPvG2qc,10789
 model_compression_toolkit/qat/pytorch/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
-model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=TaciVmT0tQhvfpp7ASxPo-feZWlUNLg4IVvx8Qpe5jA,12963
+model_compression_toolkit/qat/pytorch/quantization_facade.py,sha256=cj1AieM-v1HZcIeNfNDX7AQzQOwUw4ZuGWw2pfuY6Ig,13230
 model_compression_toolkit/qat/pytorch/quantizer/__init__.py,sha256=xYa4C8pr9cG1f3mQQcBXO_u3IdJN-zl7leZxuXDs86w,1003
 model_compression_toolkit/qat/pytorch/quantizer/base_pytorch_qat_quantizer.py,sha256=FnhuFCuQoSf78FM1z1UZgXXd3k-mKSM7i9dYOuJUmeA,2213
 model_compression_toolkit/qat/pytorch/quantizer/quantization_builder.py,sha256=e8Yfqbc552iAiP4Zxbd2ht1A3moRFGnV_KRGDm9Gw_g,5709
@@ -422,7 +422,7 @@ model_compression_toolkit/target_platform_capabilities/target_platform/targetpla
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/attribute_filter.py,sha256=-riVk2KPy94nYuviaZzZPc6j5vObhD9-6fGryuSLZ9c,8759
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/current_tpc.py,sha256=GcLSXZLxtcE9SxSKdlvo10ba9mqVk_MBiwrvvjSH8H0,2046
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/layer_filter_params.py,sha256=Cl6-mACpje2jM8RJkibbqE3hvTkFR3r26-lW021mIiA,4019
-model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py,sha256=HpJ_zzYHpSMbJ5K-IDhmP-8mwCYconaK17NSIJ3R6iI,6743
+model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py,sha256=Va9-f7M2OK3kOai5AwT-wI2zuezC9is9lwq5OOMhT_4,6733
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities.py,sha256=m6p9pO_xqGcp-0jAVRaOJww67oSQ6gChCD45_W833Gw,9819
 model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/target_platform_capabilities_component.py,sha256=FvrYI0Qy7DCmDp2gyUYyCZq5pY84JgLtJqSIiVTJ8Ss,1030
 model_compression_toolkit/target_platform_capabilities/tpc_models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -472,8 +472,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=SbvRlIdE32PEBsINt1bhSqvrKL_zbM9V-aeSkOn-sw4,3083
-mct_nightly-1.11.0.20240317.91316.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-1.11.0.20240317.91316.dist-info/METADATA,sha256=PKhfKhLQsDpsv2WN0bqt3yB2heGoLsRU7Eu6wf4JFrs,18527
-mct_nightly-1.11.0.20240317.91316.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-1.11.0.20240317.91316.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-1.11.0.20240317.91316.dist-info/RECORD,,
+mct_nightly-1.11.0.20240319.407.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-1.11.0.20240319.407.dist-info/METADATA,sha256=NODMa-MyTyAJxebf2FD53plT_HC0BdlshyXRbOi9vio,18525
+mct_nightly-1.11.0.20240319.407.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-1.11.0.20240319.407.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-1.11.0.20240319.407.dist-info/RECORD,,

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__ = "1.11.0.20240317.091316"
+__version__ = "1.11.0.20240319.000407"

model_compression_toolkit/core/__init__.py

@@ -25,3 +25,5 @@ from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core.keras.kpi_data_facade import keras_kpi_data
 from model_compression_toolkit.core.pytorch.kpi_data_facade import pytorch_kpi_data
+from model_compression_toolkit.core.common.mixed_precision.distance_weighting import MpDistanceWeighting
+

model_compression_toolkit/core/common/mixed_precision/kpi_tools/kpi_data.py

@@ -47,13 +47,15 @@ def compute_kpi_data(in_model: Any,
 
     """
 
+    # We assume that the kpi_data API is used to compute the model KPI for mixed precision scenario,
+    # so we run graph preparation under the assumption of enabled mixed precision.
     transformed_graph = graph_preparation_runner(in_model,
                                                  representative_data_gen,
                                                  core_config.quantization_config,
                                                  fw_info,
                                                  fw_impl,
                                                  tpc,
-                                                 mixed_precision_enable=core_config.mixed_precision_enable)
+                                                 mixed_precision_enable=True)
 
     # Compute parameters sum
     weights_params = compute_nodes_weights_params(graph=transformed_graph, fw_info=fw_info)

model_compression_toolkit/core/common/mixed_precision/mixed_precision_quantization_config.py

@@ -16,13 +16,11 @@
 from typing import List, Callable
 
 from model_compression_toolkit.core.common.mixed_precision.distance_weighting import MpDistanceWeighting
-from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 
 
 class MixedPrecisionQuantizationConfig:
 
     def __init__(self,
-                 target_kpi: KPI = None,
                  compute_distance_fn: Callable = None,
                  distance_weighting_method: MpDistanceWeighting = MpDistanceWeighting.AVG,
                  num_of_images: int = 32,
@@ -36,7 +34,6 @@ class MixedPrecisionQuantizationConfig:
         Class with mixed precision parameters to quantize the input model.
 
         Args:
-            target_kpi (KPI): KPI to constraint the search of the mixed-precision configuration for the model.
             compute_distance_fn (Callable): Function to compute a distance between two tensors. If None, using pre-defined distance methods based on the layer type for each layer.
             distance_weighting_method (MpDistanceWeighting): MpDistanceWeighting enum value that provides a function to use when weighting the distances among different layers when computing the sensitivity metric.
             num_of_images (int): Number of images to use to evaluate the sensitivity of a mixed-precision model comparing to the float model.
@@ -49,7 +46,6 @@ class MixedPrecisionQuantizationConfig:
 
         """
 
-        self.target_kpi = target_kpi
         self.compute_distance_fn = compute_distance_fn
         self.distance_weighting_method = distance_weighting_method
         self.num_of_images = num_of_images
@@ -67,13 +63,21 @@ class MixedPrecisionQuantizationConfig:
 
         self.metric_normalization_threshold = metric_normalization_threshold
 
-    def set_target_kpi(self, target_kpi: KPI):
+        self._mixed_precision_enable = False
+
+    def set_mixed_precision_enable(self):
+        """
+        Set a flag in mixed precision config indicating that mixed precision is enabled.
         """
-        Setting target KPI in mixed precision config.
 
-        Args:
-            target_kpi: A target KPI to set.
+        self._mixed_precision_enable = True
 
+    @property
+    def mixed_precision_enable(self):
         """
+        A property that indicates whether mixed precision quantization is enabled.
 
-        self.target_kpi = target_kpi
+        Returns: True if mixed precision quantization is enabled
+
+        """
+        return self._mixed_precision_enable

model_compression_toolkit/core/common/mixed_precision/mixed_precision_search_facade.py

@@ -47,6 +47,7 @@ search_methods = {
 def search_bit_width(graph_to_search_cfg: Graph,
                      fw_info: FrameworkInfo,
                      fw_impl: FrameworkImplementation,
+                     target_kpi: KPI,
                      mp_config: MixedPrecisionQuantizationConfig,
                      representative_data_gen: Callable,
                      search_method: BitWidthSearchMethod = BitWidthSearchMethod.INTEGER_PROGRAMMING,
@@ -63,6 +64,7 @@ def search_bit_width(graph_to_search_cfg: Graph,
         graph_to_search_cfg: Graph to search a MP configuration for.
         fw_info: FrameworkInfo object about the specific framework (e.g., attributes of different layers' weights to quantize).
         fw_impl: FrameworkImplementation object with specific framework methods implementation.
+        target_kpi: Target KPI to bound our feasible solution space s.t the configuration does not violate it.
         mp_config: Mixed-precision quantization configuration.
         representative_data_gen: Dataset to use for retrieving images for the models inputs.
         search_method: BitWidthSearchMethod to define which searching method to use.
@@ -74,7 +76,6 @@ def search_bit_width(graph_to_search_cfg: Graph,
         bit-width index on the node).
 
     """
-    target_kpi = mp_config.target_kpi
 
     # target_kpi have to be passed. If it was not passed, the facade is not supposed to get here by now.
     if target_kpi is None:

model_compression_toolkit/core/common/pruning/pruning_config.py

@@ -20,14 +20,20 @@ from model_compression_toolkit.constants import PRUNING_NUM_SCORE_APPROXIMATIONS
 
 class ImportanceMetric(Enum):
     """
-    Enum for specifying the metric used to determine the importance of channels when pruning.
+    Enum for specifying the metric used to determine the importance of channels when pruning:
+
+    LFH - Label-Free Hessian uses hessian info for measuring each channel's sensitivity.
+
     """
     LFH = 0  # Score based on the Hessian matrix w.r.t. layers weights, to determine channel importance without labels.
 
 
 class ChannelsFilteringStrategy(Enum):
     """
-    Enum for specifying the strategy used for filtering (pruning) channels.
+    Enum for specifying the strategy used for filtering (pruning) channels:
+
+    GREEDY - Prune the least important channel groups up to allowed resources in the KPI (for now, only weights_memory is considered).
+
     """
     GREEDY = 0  # Greedy strategy for pruning channels based on importance metrics.
 

model_compression_toolkit/core/common/pruning/pruning_info.py

@@ -26,23 +26,16 @@ class PruningInfo:
     and importance scores for each layer. This class acts as a container for accessing
     pruning-related metadata.
 
-    Attributes:
-        pruning_masks (Dict[BaseNode, np.ndarray]): Stores the pruning masks for each layer.
-            A pruning mask is an array where each element indicates whether the corresponding
-            channel or neuron has been pruned (0) or kept (1).
-        importance_scores (Dict[BaseNode, np.ndarray]): Stores the importance scores for each layer.
-            Importance scores quantify the significance of each channel in the layer.
     """
 
     def __init__(self,
                  pruning_masks: Dict[BaseNode, np.ndarray],
                  importance_scores: Dict[BaseNode, np.ndarray]):
         """
-        Initializes the PruningInfo with pruning masks and importance scores.
-
         Args:
-            pruning_masks (Dict[BaseNode, np.ndarray]): Pruning masks for each layer.
-            importance_scores (Dict[BaseNode, np.ndarray]): Importance scores for each layer.
+            pruning_masks (Dict[BaseNode, np.ndarray]): Stores the pruning masks for each layer. A pruning mask is an array where each element indicates whether the corresponding channel or neuron has been pruned (0) or kept (1).
+            importance_scores (Dict[BaseNode, np.ndarray]): Stores the importance scores for each layer. Importance scores quantify the significance of each channel in the layer.
+
         """
         self._pruning_masks = pruning_masks
         self._importance_scores = importance_scores

model_compression_toolkit/core/common/quantization/core_config.py

@@ -30,14 +30,19 @@ class CoreConfig:
 
         Args:
             quantization_config (QuantizationConfig): Config for quantization.
-            mixed_precision_config (MixedPrecisionQuantizationConfig): Config for mixed precision quantization (optional, default=None).
+            mixed_precision_config (MixedPrecisionQuantizationConfig): Config for mixed precision quantization.
+            If None, a default MixedPrecisionQuantizationConfig is used.
             debug_config (DebugConfig): Config for debugging and editing the network quantization process.
         """
         self.quantization_config = quantization_config
-        self.mixed_precision_config = mixed_precision_config
         self.debug_config = debug_config
 
+        if mixed_precision_config is None:
+            self.mixed_precision_config = MixedPrecisionQuantizationConfig()
+        else:
+            self.mixed_precision_config = mixed_precision_config
+
     @property
     def mixed_precision_enable(self):
-        return self.mixed_precision_config is not None
+        return self.mixed_precision_config is not None and self.mixed_precision_config.mixed_precision_enable
 

model_compression_toolkit/core/common/quantization/set_node_quantization_config.py

@@ -71,7 +71,7 @@ def set_quantization_configs_to_node(node: BaseNode,
         quant_config: Quantization configuration to generate the node's configurations from.
         fw_info: Information needed for quantization about the specific framework.
         tpc: TargetPlatformCapabilities to get default OpQuantizationConfig.
-        mixed_precision_enable: is mixed precision enabled
+        mixed_precision_enable: is mixed precision enabled.
     """
     node_qc_options = node.get_qco(tpc)
 

model_compression_toolkit/core/graph_prep_runner.py

@@ -57,7 +57,8 @@ def graph_preparation_runner(in_model: Any,
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         tpc: TargetPlatformCapabilities object that models the inference target platform and
             the attached framework operator's information.
-        tb_w: TensorboardWriter object for logging
+        tb_w: TensorboardWriter object for logging.
+        mixed_precision_enable: is mixed precision enabled.
 
     Returns:
         An internal graph representation of the input model.
@@ -103,7 +104,7 @@ def get_finalized_graph(initial_graph: Graph,
             kernel channels indices, groups of layers by how they should be quantized, etc.)
         tb_w (TensorboardWriter): TensorboardWriter object to use for logging events such as graphs, histograms, etc.
         fw_impl (FrameworkImplementation): FrameworkImplementation object with a specific framework methods implementation.
-            mixed_precision_enable: is mixed precision enabled.
+        mixed_precision_enable: is mixed precision enabled.
 
     Returns: Graph object that represents the model, after applying all required modifications to it.
     """

model_compression_toolkit/core/pytorch/kpi_data_facade.py

@@ -38,7 +38,7 @@ if FOUND_TORCH:
 
     def pytorch_kpi_data(in_model: Module,
                          representative_data_gen: Callable,
-                         core_config: CoreConfig = CoreConfig(mixed_precision_config=MixedPrecisionQuantizationConfig()),
+                         core_config: CoreConfig = CoreConfig(),
                          target_platform_capabilities: TargetPlatformCapabilities = PYTORCH_DEFAULT_TPC) -> KPI:
         """
         Computes KPI data that can be used to calculate the desired target KPI for mixed-precision quantization.

model_compression_toolkit/core/runner.py

@@ -47,6 +47,7 @@ def core_runner(in_model: Any,
                 fw_info: FrameworkInfo,
                 fw_impl: FrameworkImplementation,
                 tpc: TargetPlatformCapabilities,
+                target_kpi: KPI = None,
                 tb_w: TensorboardWriter = None):
     """
     Quantize a trained model using post-training quantization.
@@ -66,6 +67,7 @@ def core_runner(in_model: Any,
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         tpc: TargetPlatformCapabilities object that models the inference target platform and
                                               the attached framework operator's information.
+        target_kpi: KPI to constraint the search of the mixed-precision configuration for the model.
         tb_w: TensorboardWriter object for logging
 
     Returns:
@@ -81,6 +83,13 @@ def core_runner(in_model: Any,
         Logger.warning('representative_data_gen generates a batch size of 1 which can be slow for optimization:'
                        ' consider increasing the batch size')
 
+    # Checking whether to run mixed precision quantization
+    if target_kpi is not None:
+        if core_config.mixed_precision_config is None:
+            Logger.critical("Provided an initialized target_kpi, that means that mixed precision quantization is "
+                            "enabled, but the provided MixedPrecisionQuantizationConfig is None.")
+        core_config.mixed_precision_config.set_mixed_precision_enable()
+
     graph = graph_preparation_runner(in_model,
                                      representative_data_gen,
                                      core_config.quantization_config,
@@ -105,13 +114,12 @@ def core_runner(in_model: Any,
     # Finalize bit widths
     ######################################
     if core_config.mixed_precision_enable:
-        if core_config.mixed_precision_config.target_kpi is None:
-            Logger.critical(f"Trying to run Mixed Precision quantization without providing a valid target KPI.")
         if core_config.mixed_precision_config.configuration_overwrite is None:
 
             bit_widths_config = search_bit_width(tg,
                                                  fw_info,
                                                  fw_impl,
+                                                 target_kpi,
                                                  core_config.mixed_precision_config,
                                                  representative_data_gen,
                                                  hessian_info_service=hessian_info_service)

model_compression_toolkit/data_generation/__init__.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
 from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TF
+from model_compression_toolkit.data_generation.common.data_generation_config import DataGenerationConfig
+from model_compression_toolkit.data_generation.common.enums import ImageGranularity, DataInitType, SchedulerType, BNLayerWeightingType, OutputLossType, BatchNormAlignemntLossType, ImagePipelineType, ImageNormalizationType
 
 if FOUND_TF:
     from model_compression_toolkit.data_generation.keras.keras_data_generation import (

model_compression_toolkit/data_generation/common/enums.py

@@ -29,10 +29,14 @@ class EnumBaseClass(Enum):
 
 class ImageGranularity(EnumBaseClass):
     """
-    An enum for choosing the image dependence granularity when generating images.
-    0. ImageWise
-    1. BatchWise
-    2. AllImages
+    An enum for choosing the image dependence granularity when generating images:
+
+    ImageWise
+
+    BatchWise
+
+    AllImages
+
     """
 
     ImageWise = 0
@@ -42,9 +46,12 @@ class ImageGranularity(EnumBaseClass):
 
 class DataInitType(EnumBaseClass):
     """
-    An enum for choosing the image dependence granularity when generating images.
-    0. Gaussian
-    1. Diverse
+    An enum for choosing the image dependence granularity when generating images:
+
+    Gaussian
+
+    Diverse
+
     """
     Gaussian = 0
     Diverse = 1
@@ -52,9 +59,14 @@ class DataInitType(EnumBaseClass):
 
 class ImagePipelineType(EnumBaseClass):
     """
-    An enum for choosing the image pipeline type for image manipulation.
-    RANDOM_CROP_FLIP: Crop and flip the images.
-    IDENTITY: Do not apply any manipulation (identity transformation).
+    An enum for choosing the image pipeline type for image manipulation:
+
+    RANDOM_CROP - Crop the images.
+
+    RANDOM_CROP_FLIP - Crop and flip the images.
+
+    IDENTITY - Do not apply any manipulation (identity transformation).
+
     """
     RANDOM_CROP = 'random_crop'
     RANDOM_CROP_FLIP = 'random_crop_flip'
@@ -63,10 +75,14 @@ class ImagePipelineType(EnumBaseClass):
 
 class ImageNormalizationType(EnumBaseClass):
     """
-    An enum for choosing the image normalization type.
-    TORCHVISION: Normalize the images using torchvision normalization.
-    KERAS_APPLICATIONS: Normalize the images using keras_applications imagenet normalization.
-    NO_NORMALIZATION: Do not apply any normalization.
+    An enum for choosing the image normalization type:
+
+    TORCHVISION - Normalize the images using torchvision normalization.
+
+    KERAS_APPLICATIONS - Normalize the images using keras_applications imagenet normalization.
+
+    NO_NORMALIZATION - Do not apply any normalization.
+
     """
     TORCHVISION = 'torchvision'
     KERAS_APPLICATIONS = 'keras_applications'
@@ -75,10 +91,14 @@ class ImageNormalizationType(EnumBaseClass):
 
 class BNLayerWeightingType(EnumBaseClass):
     """
-   An enum for choosing the layer weighting type.
-   AVERAGE: Use the same weight per layer.
-   FIRST_LAYER_MULTIPLIER: Use a multiplier for the first layer, all other layers with the same weight.
-   GRAD: Use gradient-based layer weighting.
+   An enum for choosing the layer weighting type:
+
+   AVERAGE - Use the same weight per layer.
+
+   FIRST_LAYER_MULTIPLIER - Use a multiplier for the first layer, all other layers with the same weight.
+
+   GRAD - Use gradient-based layer weighting.
+
    """
     AVERAGE = 'average'
     FIRST_LAYER_MULTIPLIER = 'first_layer_multiplier'
@@ -87,18 +107,24 @@ class BNLayerWeightingType(EnumBaseClass):
 
 class BatchNormAlignemntLossType(EnumBaseClass):
     """
-    An enum for choosing the BatchNorm alignment loss type.
-    L2_SQUARE: Use L2 square loss for BatchNorm alignment.
+    An enum for choosing the BatchNorm alignment loss type:
+
+    L2_SQUARE - Use L2 square loss for BatchNorm alignment.
+
     """
     L2_SQUARE = 'l2_square'
 
 
 class OutputLossType(EnumBaseClass):
     """
-    An enum for choosing the output loss type.
-    NONE: No output loss is applied.
-    MIN_MAX_DIFF: Use min-max difference as the output loss.
-    REGULARIZED_MIN_MAX_DIFF: Use regularized min-max difference as the output loss.
+    An enum for choosing the output loss type:
+
+    NONE - No output loss is applied.
+
+    MIN_MAX_DIFF - Use min-max difference as the output loss.
+
+    REGULARIZED_MIN_MAX_DIFF - Use regularized min-max difference as the output loss.
+
     """
     NONE = 'none'
     MIN_MAX_DIFF = 'min_max_diff'
@@ -107,9 +133,12 @@ class OutputLossType(EnumBaseClass):
 
 class SchedulerType(EnumBaseClass):
     """
-    An enum for choosing the scheduler type for the optimizer.
-    REDUCE_ON_PLATEAU: Use the ReduceOnPlateau scheduler.
-    STEP: Use the Step scheduler.
+    An enum for choosing the scheduler type for the optimizer:
+
+    REDUCE_ON_PLATEAU - Use the ReduceOnPlateau scheduler.
+
+    STEP - Use the Step scheduler.
+
     """
     REDUCE_ON_PLATEAU = 'reduce_on_plateau'
     STEP = 'step'

model_compression_toolkit/data_generation/keras/keras_data_generation.py

@@ -131,7 +131,36 @@ if FOUND_TF:
 
         Returns:
             List[tf.Tensor]: Finalized list containing generated images.
+
+        Examples:
+
+            In this example, we'll walk through generating images using a simple Keras model and a data generation configuration. The process involves creating a model, setting up a data generation configuration, and finally generating images with specified parameters.
+
+            Start by importing the Model Compression Toolkit (MCT), TensorFlow, and some layers from `tensorflow.keras`:
+
+            >>> import model_compression_toolkit as mct
+            >>> from tensorflow.keras.models import Sequential
+            >>> from tensorflow.keras.layers import Conv2D, BatchNormalization, Flatten, Dense, Reshape
+
+            Next, define a simple Keras model:
+
+            >>> model = Sequential([Conv2D(2, 3, input_shape=(8,8,3)), BatchNormalization(), Flatten(), Dense(10)])
+
+            Configure the data generation process using `get_keras_data_generation_config`. This function allows customization of the data generation process. For simplicity, this example sets the number of iterations (`n_iter`) to 1 and the batch size (`data_gen_batch_size`) to 2.
+
+            >>> config = mct.data_generation.get_keras_data_generation_config(n_iter=1, data_gen_batch_size=2)
+
+            Finally, use the `keras_data_generation_experimental` function to generate images based on the model and data generation configuration.
+            Notice that this function is experimental and may change in future versions of MCT.
+            The `n_images` parameter specifies the number of images to generate, and `output_image_size` sets the size of the generated images.
+
+            >>> generated_images = mct.data_generation.keras_data_generation_experimental(model=model, n_images=4, output_image_size=(8, 8), data_generation_config=config)
+
+            The generated images can then be used for various purposes, such as data-free quantization.
+
+
         """
+
         Logger.warning(f"keras_data_generation_experimental is experimental "
                        f"and is subject to future changes."
                        f"If you encounter an issue, please open an issue in our GitHub "

model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py

@@ -129,7 +129,7 @@ if FOUND_TORCH:
     def pytorch_data_generation_experimental(
             model: Module,
             n_images: int,
-            output_image_size: Tuple,
+            output_image_size: int,
             data_generation_config: DataGenerationConfig) -> List[Tensor]:
         """
         Function to perform data generation using the provided model and data generation configuration.
@@ -137,11 +137,38 @@ if FOUND_TORCH:
         Args:
             model (Module): PyTorch model to generate data for.
             n_images (int): Number of images to generate.
-            output_image_size (Tuple): Size of the output images.
+            output_image_size (int): The hight and width size of the output images.
             data_generation_config (DataGenerationConfig): Configuration for data generation.
 
         Returns:
             List[Tensor]: Finalized list containing generated images.
+
+        Examples:
+
+            In this example, we'll walk through generating images using a simple PyTorch model and a data generation configuration. The process involves creating a model, setting up a data generation configuration, and finally generating images with specified parameters.
+
+            Start by importing the Model Compression Toolkit (MCT), PyTorch, and some modules from `torch.nn`:
+
+            >>> import model_compression_toolkit as mct
+            >>> import torch.nn as nn
+            >>> from torch.nn import Conv2d, BatchNorm2d, Flatten, Linear
+
+            Next, define a simple PyTorch model:
+
+            >>> model = nn.Sequential(nn.Conv2d(3, 2, 3), nn.BatchNorm2d(2), nn.Flatten(), nn.Linear(2*6*6, 10))
+
+            Configure the data generation process using `get_pytorch_data_generation_config`. This function allows customization of the data generation process. For simplicity, this example sets the number of iterations (`n_iter`) to 1 and the batch size (`data_gen_batch_size`) to 2.
+
+            >>> config = mct.data_generation.get_pytorch_data_generation_config(n_iter=1, data_gen_batch_size=2)
+
+            Finally, use the `pytorch_data_generation_experimental` function to generate images based on the model and data generation configuration.
+            Notice that this function is experimental and may change in future versions of MCT.
+            The `n_images` parameter specifies the number of images to generate, and `output_image_size` sets the size of the generated images.
+
+            >>> generated_images = mct.data_generation.pytorch_data_generation_experimental(model=model, n_images=4, output_image_size=8, data_generation_config=config)
+
+            The generated images can then be used for various purposes, such as data-free quantization.
+
         """
 
         Logger.warning(f"pytorch_data_generation_experimental is experimental "

model_compression_toolkit/exporter/model_exporter/fw_agonstic/quantization_format.py

@@ -16,6 +16,16 @@ from enum import Enum
 
 
 class QuantizationFormat(Enum):
+    """
+    Specify which quantization format to use for exporting a quantized model.
+
+    FAKELY_QUANT - Weights and activations are quantized but represented using float data type.
+
+    INT8 - Weights and activations are represented using 8-bit integer data type.
+
+    MCTQ - Weights and activations are quantized using mct_quantizers custom quantizers.
+
+    """
     FAKELY_QUANT = 0
     INT8 = 1
     MCTQ = 2

model_compression_toolkit/exporter/model_exporter/keras/keras_export_facade.py

@@ -42,21 +42,17 @@ if FOUND_TF:
                            serialization_format: KerasExportSerializationFormat = KerasExportSerializationFormat.KERAS,
                            quantization_format : QuantizationFormat = QuantizationFormat.MCTQ) -> Dict[str, type]:
         """
-        Export a Keras quantized model to a h5 or tflite model.
+        Export a Keras quantized model to a .keras or .tflite format model (according to serialization_format).
         The model will be saved to the path in save_model_path.
-        keras_export_model supports the combination of QuantizationFormat.FAKELY_QUANT (where weights
-        and activations are float fakely-quantized values) and KerasExportSerializationFormat.KERAS_H5 (where the model
-        will be saved to h5 model) or the combination of KerasExportSerializationFormat.TFLITE (where the model will be
-        saved to tflite model) with QuantizationFormat.FAKELY_QUANT or QuantizationFormat.INT8 (where weights and
-        activations are represented using 8bits integers).
+        Models that are exported to .keras format can use quantization_format of QuantizationFormat.MCTQ or QuantizationFormat.FAKELY_QUANT.
+        Models that are exported to .tflite format can use quantization_format of QuantizationFormat.INT8 or QuantizationFormat.FAKELY_QUANT.
 
         Args:
             model: Model to export.
             save_model_path: Path to save the model.
             is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
-            serialization_format: Format to export the model according to (by default
-            KerasExportSerializationFormat.KERAS_H5).
-            quantization_format: Format of how quantizers are exported (fakely-quant, int8, MCTQ quantizers).
+            serialization_format: Format to export the model according to (KerasExportSerializationFormat.KERAS, by default).
+            quantization_format: Format of how quantizers are exported (MCTQ quantizers, by default).
 
         Returns:
             Custom objects dictionary needed to load the model.

model_compression_toolkit/gptq/common/gptq_config.py

@@ -19,9 +19,12 @@ from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT
 
 class RoundingType(Enum):
     """
-    An enum for choosing the GPTQ rounding methods
-    0. STRAIGHT-THROUGH ESTIMATOR
-    1. SoftQuantizer
+    An enum for choosing the GPTQ rounding methods:
+
+    STE - STRAIGHT-THROUGH ESTIMATOR
+
+    SoftQuantizer - SoftQuantizer
+
     """
     STE = 0
     SoftQuantizer = 1

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -116,6 +116,7 @@ if FOUND_TF:
     def keras_gradient_post_training_quantization(in_model: Model, representative_data_gen: Callable,
                                                   gptq_config: GradientPTQConfig,
                                                   gptq_representative_data_gen: Callable = None,
+                                                  target_kpi: KPI = None,
                                                   core_config: CoreConfig = CoreConfig(),
                                                   target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, UserInformation]:
         """
@@ -139,6 +140,7 @@ if FOUND_TF:
             representative_data_gen (Callable): Dataset used for calibration.
             gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
+            target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
@@ -166,6 +168,12 @@ if FOUND_TF:
 
             >>> config = mct.core.CoreConfig()
 
+            If mixed precision is desired, create an MCT core config with a mixed-precision configuration, to quantize a model
+            with different bitwidths for different layers.
+            The candidates bitwidth for quantization should be defined in the target platform model:
+
+            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=1))
+
             For mixed-precision set a target KPI object:
             Create a KPI object to limit our returned model's size. Note that this value affects only coefficients
             that should be quantized (for example, the kernel of Conv2D in Keras will be affected by this value,
@@ -173,19 +181,13 @@ if FOUND_TF:
 
             >>> kpi = mct.core.KPI(model.count_params() * 0.75)  # About 0.75 of the model size when quantized with 8 bits.
 
-            If mixed precision is desired, create an MCT core config with a mixed-precision configuration, to quantize a model
-            with different bitwidths for different layers.
-            The candidates bitwidth for quantization should be defined in the target platform model:
-
-            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=1, target_kpi=kpi))
-
             Create GPTQ config:
 
             >>> gptq_config = mct.gptq.get_keras_gptq_config(n_epochs=1)
 
             Pass the model with the representative dataset generator to get a quantized model:
 
-            >>> quantized_model, quantization_info = mct.gptq.keras_gradient_post_training_quantization(model, repr_datagen, gptq_config, core_config=config)
+            >>> quantized_model, quantization_info = mct.gptq.keras_gradient_post_training_quantization(model, repr_datagen, gptq_config, target_kpi=kpi, core_config=config)
 
         """
         KerasModelValidation(model=in_model,

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -94,6 +94,7 @@ if FOUND_TORCH:
 
     def pytorch_gradient_post_training_quantization(model: Module,
                                                     representative_data_gen: Callable,
+                                                    target_kpi: KPI = None,
                                                     core_config: CoreConfig = CoreConfig(),
                                                     gptq_config: GradientPTQConfig = None,
                                                     gptq_representative_data_gen: Callable = None,
@@ -117,6 +118,7 @@ if FOUND_TORCH:
         Args:
             model (Module): Pytorch model to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
+            target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
             gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
@@ -174,6 +176,7 @@ if FOUND_TORCH:
                                                                      fw_info=DEFAULT_PYTORCH_INFO,
                                                                      fw_impl=fw_impl,
                                                                      tpc=target_platform_capabilities,
+                                                                     target_kpi=target_kpi,
                                                                      tb_w=tb_w)
 
         # ---------------------- #

model_compression_toolkit/pruning/keras/pruning_facade.py

@@ -40,8 +40,7 @@ if FOUND_TF:
                                    target_kpi: KPI,
                                    representative_data_gen: Callable,
                                    pruning_config: PruningConfig = PruningConfig(),
-                                   target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> \
-            Tuple[Model, PruningInfo]:
+                                   target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, PruningInfo]:
         """
         Perform structured pruning on a Keras model to meet a specified target KPI.
         This function prunes the provided model according to the target KPI by grouping and pruning
@@ -59,12 +58,14 @@ if FOUND_TF:
             target_kpi (KPI): The target Key Performance Indicators to be achieved through pruning.
             representative_data_gen (Callable): A function to generate representative data for pruning analysis.
             pruning_config (PruningConfig): Configuration settings for the pruning process. Defaults to standard config.
-            target_platform_capabilities (TargetPlatformCapabilities): Platform-specific constraints and capabilities.
-                Defaults to DEFAULT_KERAS_TPC.
+            target_platform_capabilities (TargetPlatformCapabilities): Platform-specific constraints and capabilities. Defaults to DEFAULT_KERAS_TPC.
 
         Returns:
             Tuple[Model, PruningInfo]: A tuple containing the pruned Keras model and associated pruning information.
 
+        Note:
+            The pruned model should be fine-tuned or retrained to recover or improve its performance post-pruning.
+
         Examples:
 
             Import MCT:

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -42,6 +42,7 @@ if FOUND_TF:
 
     def keras_post_training_quantization(in_model: Model,
                                          representative_data_gen: Callable,
+                                         target_kpi: KPI = None,
                                          core_config: CoreConfig = CoreConfig(),
                                          target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC):
         """
@@ -60,6 +61,7 @@ if FOUND_TF:
          Args:
              in_model (Model): Keras model to quantize.
              representative_data_gen (Callable): Dataset used for calibration.
+             target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
              target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
@@ -135,6 +137,7 @@ if FOUND_TF:
                                                fw_info=fw_info,
                                                fw_impl=fw_impl,
                                                tpc=target_platform_capabilities,
+                                               target_kpi=target_kpi,
                                                tb_w=tb_w)
 
         tg = ptq_runner(tg, representative_data_gen, core_config, fw_info, fw_impl, tb_w)

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -41,6 +41,7 @@ if FOUND_TORCH:
 
     def pytorch_post_training_quantization(in_module: Module,
                                            representative_data_gen: Callable,
+                                           target_kpi: KPI = None,
                                            core_config: CoreConfig = CoreConfig(),
                                            target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
         """
@@ -59,6 +60,7 @@ if FOUND_TORCH:
         Args:
             in_module (Module): Pytorch module to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
+            target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to.
 
@@ -107,6 +109,7 @@ if FOUND_TORCH:
                                                fw_info=DEFAULT_PYTORCH_INFO,
                                                fw_impl=fw_impl,
                                                tpc=target_platform_capabilities,
+                                               target_kpi=target_kpi,
                                                tb_w=tb_w)
 
         tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_PYTORCH_INFO, fw_impl, tb_w)

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -87,6 +87,7 @@ if FOUND_TF:
 
     def keras_quantization_aware_training_init_experimental(in_model: Model,
                                                             representative_data_gen: Callable,
+                                                            target_kpi: KPI = None,
                                                             core_config: CoreConfig = CoreConfig(),
                                                             qat_config: QATConfig = QATConfig(),
                                                             target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC):
@@ -108,6 +109,7 @@ if FOUND_TF:
          Args:
              in_model (Model): Keras model to quantize.
              representative_data_gen (Callable): Dataset used for initial calibration.
+             target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
              qat_config (QATConfig): QAT configuration
              target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
@@ -157,7 +159,7 @@ if FOUND_TF:
              Pass the model, the representative dataset generator, the configuration and the target KPI to get a
              quantized model:
 
-             >>> quantized_model, quantization_info, custom_objects = mct.qat.keras_quantization_aware_training_init_experimental(model, repr_datagen, kpi, core_config=core_config)
+             >>> quantized_model, quantization_info, custom_objects = mct.qat.keras_quantization_aware_training_init_experimental(model, repr_datagen, kpi, core_config=config)
 
              Use the quantized model for fine-tuning. For loading the model from file, use the custom_objects dictionary:
 
@@ -191,6 +193,7 @@ if FOUND_TF:
                                                fw_info=DEFAULT_KERAS_INFO,
                                                fw_impl=fw_impl,
                                                tpc=target_platform_capabilities,
+                                               target_kpi=target_kpi,
                                                tb_w=tb_w)
 
         tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_KERAS_INFO, fw_impl, tb_w)

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -75,6 +75,7 @@ if FOUND_TORCH:
 
     def pytorch_quantization_aware_training_init_experimental(in_model: Module,
                                                               representative_data_gen: Callable,
+                                                              target_kpi: KPI = None,
                                                               core_config: CoreConfig = CoreConfig(),
                                                               qat_config: QATConfig = QATConfig(),
                                                               target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
@@ -96,6 +97,7 @@ if FOUND_TORCH:
          Args:
              in_model (Model): Pytorch model to quantize.
              representative_data_gen (Callable): Dataset used for initial calibration.
+             target_kpi (KPI): KPI object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
              qat_config (QATConfig): QAT configuration
              target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Pytorch model according to.
@@ -158,6 +160,7 @@ if FOUND_TORCH:
                                                fw_info=DEFAULT_PYTORCH_INFO,
                                                fw_impl=fw_impl,
                                                tpc=target_platform_capabilities,
+                                               target_kpi=target_kpi,
                                                tb_w=tb_w)
 
         tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_PYTORCH_INFO, fw_impl, tb_w)

model_compression_toolkit/target_platform_capabilities/target_platform/targetplatform2framework/operations_to_layers.py

@@ -35,9 +35,9 @@ class OperationsSetToLayers(TargetPlatformCapabilitiesComponent):
         Args:
             op_set_name (str): Name of OperatorsSet to associate with layers.
             layers (List[Any]): List of layers/FilterLayerParams to associate with OperatorsSet.
-            attr_mapping (dict): A mapping between a general attribute name to a DefaultDict that maps a layer
-+                type to the layer's framework name of this attribute (the dictionary type is not specified to
-+                handle circular dependency).
+            attr_mapping (dict): A mapping between a general attribute name to a DefaultDict that maps a layer type
+            to the layer's framework name of this attribute (the dictionary type is not specified to handle circular
+            dependency).
         """
         self.layers = layers
         self.attr_mapping = attr_mapping