changed replace_operations method

b_asic/signal_flow_graph.py

@@ -414,12 +414,12 @@ class SFG(AbstractOperation):
         operation: The operation used for replacement.
         """
 
+        _sfg = self()
+
         inputs = []
         outputs = []
 
-        for _operation in self.operations:
-            if _operation.graph_id not in operation_ids:
-                continue
+        for _operation in [_sfg.find_by_id(_id) for _id in operation_ids]:
             
             # Retrive input operations
             for _signal in _operation.input_signals:
@@ -431,8 +431,8 @@ class SFG(AbstractOperation):
                 if _signal.destination.operation.graph_id not in operation_ids:
                     outputs.append(_signal.destination.operation)
 
-        assert len(inputs) + len(outputs) == \
-            operation.input_count + operation.output_count, "The input and output count must match"
+        assert len(inputs) == operation.input_count, "The input count must match"
+        assert len(outputs) == operation.output_count, "The output count must match"
 
         for index_in, _input in enumerate(inputs):
             for _signal in _input.output_signals:
@@ -443,8 +443,8 @@ class SFG(AbstractOperation):
             for _signal in _output.input_signals:
                 _signal.remove_source()
                 _signal.set_source(operation.output(index_out))
-        
-        return self()
+       
+        return _sfg
 
     def _evaluate_source(self, src: OutputPort, results: MutableResultMap, registers: MutableRegisterMap, prefix: str) -> Number:
         src_prefix = prefix