Fix RemoveIdentityOps not correctly handling ops following fork-nodes

Author: iksnagreb

src/qonnx/transformation/remove.py

@@ -25,9 +25,8 @@
 # CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-
 import numpy as np
+import warnings
 
 from qonnx.core.modelwrapper import ModelWrapper
 from qonnx.transformation.base import Transformation
@@ -58,21 +57,45 @@ def apply(self, model: ModelWrapper):
 
 
 def remove_node_and_rewire(model, node):
+    # Currently cannot remove and rewire join-nodes, probably not necessary to
+    # support this
+    if model.is_join_node(node):
+        # Log this as a warning, so the user is aware of this, there might be
+        # somthing wrong or some checks missing at the caller site
+        warnings.warn(
+            "Tried to remove join-node operation: Currently not supported"
+        )
+        # Exit the function here without doing anything
+        return
+    # We already know that node is not a join-node, thus to rewire, we only need
+    # to check the single producer
     producer = model.find_producer(node.input[0])
-    if producer is not None:
-        # wire output tensor to
-        # output of producer node
+    # If there is a producer which is not a fork-node, rewiring is simple
+    if producer is not None and not model.is_fork_node(producer):
+        # Rewire by skipping the node, letting the producer directly feed the
+        # nodes output.
+        #   TODO: Check whether this already covers fork-node identities?
         producer.output[0] = node.output[0]
+    # If there is no producer or the producer forks, rewiring is a bit more
+    # complicated
     else:
-        # node is first in graph
+        # Now it depends on the successor nodes to rewire their inputs
         successors = model.find_direct_successors(node)
+        # Singular node detached from the rest of the graph?
         assert successors is not None, "Whole graph is one node."
-        for succ in successors:
-            for i, s_inp in enumerate(succ.input):
+        # We need to rewire the input of each successor to not detach parts of
+        # the graph
+        for successor in successors:
+            # Find the inputs of the successor which are produced by the node to
+            # be removed
+            for i, s_inp in enumerate(successor.input):
+                # Note: This might happen multiple times?
                 if s_inp == node.output[0]:
-                    # rewire successor's input directly to graph input
-                    succ.input[i] = node.input[0]
-    # remove node
+                    # Rewire successor's input directly to nodes input
+                    # Note: Node may not be a join-node, but there is probably
+                    # no such thing as join-node identity anyway
+                    successor.input[i] = node.input[0]
+    # Remove node
     model.graph.node.remove(node)