mct-nightly 2.0.0.20240421.444__py3-none-any.whl → 2.0.0.20240423.418__py3-none-any.whl

{mct_nightly-2.0.0.20240421.444.dist-info → mct_nightly-2.0.0.20240423.418.dist-info}/METADATA

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

{mct_nightly-2.0.0.20240421.444.dist-info → mct_nightly-2.0.0.20240423.418.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=Sp-F3_qSg46tG1NVEA0zcJ4p7BV5WO8mh4387kEYV18,1573
+model_compression_toolkit/__init__.py,sha256=5k3NcM1GZLJ0pv5dLUA6lhAdMxITH-pJzNUew1byBi4,1573
 model_compression_toolkit/constants.py,sha256=yIJyJ-e1WrDeKD9kG15qkqfYnoj7J1J2CxnJDt008ik,3756
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -135,7 +135,7 @@ model_compression_toolkit/core/common/substitutions/batchnorm_reconstruction.py,
 model_compression_toolkit/core/common/substitutions/batchnorm_refusing.py,sha256=YqLKiO5gFBEvI6noAWeMME1JHaYUaGFMglVFg8AqGjc,10028
 model_compression_toolkit/core/common/substitutions/linear_collapsing.py,sha256=iEtzbWCDXP6EDkTZCtREQ0rpMxhQ2kM9zlcP_0KLq9I,12367
 model_compression_toolkit/core/common/substitutions/linear_collapsing_substitution.py,sha256=uoauhmncQqUBNvD-qCLIXsIbl_IzrbxSKdxiMig-5W4,2406
-model_compression_toolkit/core/common/substitutions/remove_identity.py,sha256=LjkedR5fnXy4LCEQ7rnVTBI-cTkdDxXtufge5Llj2J0,2038
+model_compression_toolkit/core/common/substitutions/remove_identity.py,sha256=VUWjc9Wo_nwSX3JTyAtrG5mprV90DwaQopAuvND30nQ,2353
 model_compression_toolkit/core/common/substitutions/residual_collapsing.py,sha256=doErjlMq-uSObYMSjA6IywSHb3Hz3QCc0HKU68ccrQ4,4767
 model_compression_toolkit/core/common/substitutions/scale_equalization.py,sha256=p57u25qdW2pimxzGwgMXEBV4S-LzXuTVAlIM7830WfU,10966
 model_compression_toolkit/core/common/substitutions/shift_negative_activation.py,sha256=cyy4qnlD-v1Gou62oHNDsf1hWLWkYfcjVv1otFrUltY,29865
@@ -483,8 +483,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=7bbzqJN8ZAycVDvZr_5xC-niTAR5df8f03Kooev_pfg,3047
-mct_nightly-2.0.0.20240421.444.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.0.0.20240421.444.dist-info/METADATA,sha256=r6IbtuZgkT0Pk8_2MxhQUriCs4oZAHFv1RM7rKvUt5k,18795
-mct_nightly-2.0.0.20240421.444.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.0.0.20240421.444.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.0.0.20240421.444.dist-info/RECORD,,
+mct_nightly-2.0.0.20240423.418.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.0.0.20240423.418.dist-info/METADATA,sha256=_EZ2ANty4q8kk8Bm6pd5P7s45Grq2qeePcVLxkeftho,18795
+mct_nightly-2.0.0.20240423.418.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.0.0.20240423.418.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.0.0.20240423.418.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__ = "2.0.0.20240421.000444"
+__version__ = "2.0.0.20240423.000418"

model_compression_toolkit/core/common/substitutions/remove_identity.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.core.common.graph.base_graph import Graph
+from model_compression_toolkit.core.common.graph.base_graph import Graph, OutTensor
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 
 
@@ -33,16 +33,24 @@ def remove_identity_node(graph: Graph,
     """
     # Retrieve the predecessor nodes of the identity node.
     prev_identity_nodes = graph.get_prev_nodes(node)
+
     # Ensure there is exactly one predecessor; otherwise, do nothing.
     if len(prev_identity_nodes) != 1:
         return graph
 
+    graph_outputs = graph.get_outputs()
+    for i, g_out in enumerate(graph_outputs):
+        if g_out.node == node:
+            graph_outputs[i] = OutTensor(node=prev_identity_nodes[0], node_out_index=g_out.node_out_index)
+
     # Reconnect the output edges of the identity node to its predecessor,
     # effectively bypassing the identity node.
     graph.reconnect_out_edges(current_node=node, new_node=prev_identity_nodes[0])
     # Remove the edge from the predecessor to the identity node.
     graph.remove_edge(prev_identity_nodes[0], node)
     # Remove the identity node from the graph.
-    graph.remove_node(node_to_remove=node)
+    graph.remove_node(node_to_remove=node,
+                      new_graph_outputs=graph_outputs
+                      )
 
     return graph