Repository: PennyLaneAI/pennylane
Pull request #2124 (2022-01-25 16:08:01)

`WireCut` nodes can be replaced with `MeasureNode` and `PrepareNode`
**Context:**
When constructing a circuit a `qml.WireCut` operator can be used to indicate a cut location. The actual circuit cutting converts this circuit into a computational graph with operators and observables as the nodes. The `qml.WireCut` operator is converted to  `MeasureNode` and `PrepareNode` which, at this stage, are also placeholders to indicate where to perform physical measurements and state preparations corresponding to a full orthonormal basis. 

**Description of the Change:**
A method to replace nodes containing `qml.WireCut` operators with `qcut.MeasureNode`s and `qcut.PrepareNode`s has been added  

**Benefits:**
Facilitates the wire cutting compiler


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2124](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (3b13808) into [sc-13390-a-user-can-convert-a-tape-to-an-accurate](https://codecov.io/gh/PennyLaneAI/pennylane/commit/d054a6b99a32255fc31136940fed51b0f41a1409?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (d054a6b) will **increase** coverage by `59.61%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@                                  Coverage Diff                                   @@
##           sc-13390-a-user-can-convert-a-tape-to-an-accurate    #2124       +/-   ##
======================================================================================
+ Coverage                                              37.47%   97.08%   +59.61%     
======================================================================================
  Files                                                    224      229        +5     
  Lines                                                  17042    17563      +521     
======================================================================================
+ Hits                                                    6386    17051    +10665     
+ Misses                                                 10656      512    -10144     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/qcut.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvcWN1dC5weQ==) | `100.00% <100.00%> (+100.00%)` | :arrow_up: |
| [pennylane/transforms/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvX19pbml0X18ucHk=) | `100.00% <0.00%> (ø)` | |
| [qchem/pennylane\_qchem/qchem/openfermion.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cWNoZW0vcGVubnlsYW5lX3FjaGVtL3FjaGVtL29wZW5mZXJtaW9uLnB5) | `100.00% <0.00%> (ø)` | |
| [qchem/pennylane\_qchem/qchem/obs.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cWNoZW0vcGVubnlsYW5lX3FjaGVtL3FjaGVtL29icy5weQ==) | `100.00% <0.00%> (ø)` | |
| [qchem/pennylane\_qchem/qchem/structure.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cWNoZW0vcGVubnlsYW5lX3FjaGVtL3FjaGVtL3N0cnVjdHVyZS5weQ==) | `97.65% <0.00%> (ø)` | |
| [qchem/pennylane\_qchem/\_version.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cWNoZW0vcGVubnlsYW5lX3FjaGVtL192ZXJzaW9uLnB5) | `100.00% <0.00%> (ø)` | |
| [qchem/pennylane\_qchem/qchem/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cWNoZW0vcGVubnlsYW5lX3FjaGVtL3FjaGVtL19faW5pdF9fLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/devices/default\_qubit\_autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9hdXRvZ3JhZC5weQ==) | `100.00% <0.00%> (+2.27%)` | :arrow_up: |
| [pennylane/templates/layers/random.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9sYXllcnMvcmFuZG9tLnB5) | `100.00% <0.00%> (+5.00%)` | :arrow_up: |
| [pennylane/devices/default\_mixed.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9taXhlZC5weQ==) | `100.00% <0.00%> (+6.95%)` | :arrow_up: |
| ... and [183 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [d054a6b...3b13808](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2124?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-0.21.0.md b/doc/releases/changelog-0.21.0.md
index 8df44a424..e93197b88 100644
--- a/doc/releases/changelog-0.21.0.md
+++ b/doc/releases/changelog-0.21.0.md
@@ -52,9 +52,6 @@
   A `WireCut` operator has been added for manual wire cut placement
   when constructing a QNode.
   [(#2093)](https://github.com/PennyLaneAI/pennylane/pull/2093)
-  A method for converting a quantum tape to a directed multigraph that is amenable
-  to graph partitioning algorithms for circuit cutting has been added.
-  [(#2107)](https://github.com/PennyLaneAI/pennylane/pull/2107)
 
 * The `RotosolveOptimizer` has been generalized to arbitrary frequency spectra
   in the cost function. Also note the changes in behaviour listed under *Breaking
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fc730e0ab..6c2ef3881 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,16 @@
 
 <h3>New features since last release</h3>
 
+* Continued development of the circuit-cutting compiler:
+  
+  A method for converting a quantum tape to a directed multigraph that is amenable
+  to graph partitioning algorithms for circuit cutting has been added.
+  [(#2107)](https://github.com/PennyLaneAI/pennylane/pull/2107)
+  
+  A method to replace `WireCut` nodes in a directed multigraph with `MeasureNode` 
+  and `PrepareNode` placeholders has been added.
+  [(#2124)](https://github.com/PennyLaneAI/pennylane/pull/2124)
+  
 <h3>Improvements</h3>
 
 <h3>Breaking changes</h3>
@@ -18,3 +28,4 @@
 
 This release contains contributions from (in alphabetical order):
 
+Anthony Hayes
diff --git a/pennylane/transforms/qcut.py b/pennylane/transforms/qcut.py
index 561c4258b..2f05df794 100644
--- a/pennylane/transforms/qcut.py
+++ b/pennylane/transforms/qcut.py
@@ -18,10 +18,137 @@ circuits to be distributed across multiple devices.
 
 from networkx import MultiDiGraph
 from pennylane.measure import MeasurementProcess
-from pennylane.operation import Operator, Tensor
+from pennylane.operation import Operation, Operator, Tensor
+from pennylane.ops.qubit.non_parametric_ops import WireCut
 from pennylane.tape import QuantumTape
 
 
+class MeasureNode(Operation):
+    """Placeholder node for measurement operations"""
+
+    num_wires = 1
+    grad_method = None
+
+
+class PrepareNode(Operation):
+    """Placeholder node for state preparations"""
+
+    num_wires = 1
+    grad_method = None
+
+
+def replace_wire_cut_node(node: WireCut, graph: MultiDiGraph):
+    """
+    Replace a :class:`~.WireCut` node in the graph with a :class:`~.MeasureNode`
+    and :class:`~.PrepareNode`.
+
+    Args:
+        node (WireCut): the  :class:`~.WireCut` node to be replaced with a :class:`~.MeasureNode`
+            and :class:`~.PrepareNode`
+        graph (MultiDiGraph): the graph containing the node to be replaced
+
+    **Example**
+
+    Consider the following circuit with a manually-placed wire cut:
+
+    .. code-block:: python
+
+        from pennylane.transforms import qcut
+
+        wire_cut = qml.WireCut(wires=0)
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.4, wires=0)
+            qml.apply(wire_cut)
+            qml.RY(0.5, wires=0)
+            qml.expval(qml.PauliZ(0))
+
+    We can find the circuit graph and remove the wire cut node using:
+
+    >>> graph = qcut.tape_to_graph(tape)
+    >>> qcut.replace_wire_cut_node(wire_cut, graph)
+    """
+    predecessors = graph.pred[node]
+    successors = graph.succ[node]
+
+    predecessor_on_wire = {}
+    for op, data in predecessors.items():
+        for d in data.values():
+            wire = d["wire"]
+            predecessor_on_wire[wire] = op
+
+    successor_on_wire = {}
+    for op, data in successors.items():
+        for d in data.values():
+            wire = d["wire"]
+            successor_on_wire[wire] = op
+
+    order = graph.nodes[node]["order"]
+    graph.remove_node(node)
+
+    for wire in node.wires:
+        predecessor = predecessor_on_wire.get(wire, None)
+        successor = successor_on_wire.get(wire, None)
+
+        meas = MeasureNode(wires=wire)
+        prep = PrepareNode(wires=wire)
+
+        # We are introducing a degeneracy in the order of the measure and prepare nodes
+        # here but the order can be inferred as MeasureNode always precedes
+        # the corresponding PrepareNode
+        graph.add_node(meas, order=order)
+        graph.add_node(prep, order=order)
+
+        graph.add_edge(meas, prep, wire=wire)
+
+        if predecessor is not None:
+            graph.add_edge(predecessor, meas, wire=wire)
+        if successor is not None:
+            graph.add_edge(prep, successor, wire=wire)
+
+
+def replace_wire_cut_nodes(graph: MultiDiGraph):
+    """
+    Replace each :class:`~.WireCut` node in the graph with a
+    :class:`~.MeasureNode` and :class:`~.PrepareNode`.
+
+    Args:
+        graph (MultiDiGraph): The graph containing the :class:`~.WireCut` nodes
+            to be replaced
+
+    **Example**
+
+    Consider the following circuit with manually-placed wire cuts:
+
+    .. code-block:: python
+
+        from pennylane.transforms import qcut
+
+        wire_cut_0 = qml.WireCut(wires=0)
+        wire_cut_1 = qml.WireCut(wires=1)
+        multi_wire_cut = qml.WireCut(wires=[0, 1])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.4, wires=0)
+            qml.apply(wire_cut_0)
+            qml.RY(0.5, wires=0)
+            qml.apply(wire_cut_1)
+            qml.CNOT(wires=[0, 1])
+            qml.apply(multi_wire_cut)
+            qml.RZ(0.6, wires=1)
+            qml.expval(qml.PauliZ(0))
+
+    We can find the circuit graph and remove all the wire cut nodes using:
+
+    >>> graph = qcut.tape_to_graph(tape)
+    >>> qcut.replace_wire_cut_nodes(graph)
+
+    """
+    for op in list(graph.nodes):
+        if isinstance(op, WireCut):
+            replace_wire_cut_node(op, graph)
+
+
 def _add_operator_node(graph: MultiDiGraph, op: Operator, order: int, wire_latest_node: dict):
     """
     Helper function to add operators as nodes during tape to graph conversion.


Test patch:
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index e531b73e5..4dfac069b 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -233,3 +233,270 @@ class TestTapeToGraph:
         for node_obs, exp_obs in zip(node_observables, expected_obs):
             assert node_obs.wires == exp_obs.wires
             assert node_obs.obs.name == exp_obs.obs.name
+
+
+class TestReplaceWireCut:
+    """
+    Tests the replacement of WireCut operation with MeasureNode and
+    PrepareNode
+    """
+
+    def test_single_wire_cut_replaced(self):
+        """
+        Tests that a single WireCut operator is replaced with a MeasureNode and
+        a PrepareNode with the correct order
+        """
+
+        wire_cut_num = 1
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.432, wires=0)
+            qml.RY(0.543, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(0.240, wires=0)
+            qml.RZ(0.133, wires=1)
+            qml.WireCut(wires=wire_cut_num)
+            qml.CNOT(wires=[1, 2])
+            qml.RX(0.432, wires=1)
+            qml.RY(0.543, wires=2)
+            qml.expval(qml.PauliZ(wires=[0]))
+
+        g = qcut.tape_to_graph(tape)
+        node_data = list(g.nodes(data=True))
+
+        wire_cut_order = {"order": 5}
+
+        qcut.replace_wire_cut_nodes(g)
+        new_node_data = list(g.nodes(data=True))
+        op_names = [op.name for op, order in new_node_data]
+
+        assert "WireCut" not in op_names
+        assert "MeasureNode" in op_names
+        assert "PrepareNode" in op_names
+
+        for op, order in new_node_data:
+            if op.name == "MeasureNode":
+                assert op.wires.tolist() == [wire_cut_num]
+                assert order == wire_cut_order
+            elif op.name == "PrepareNode":
+                assert op.wires.tolist() == [wire_cut_num]
+                assert order == wire_cut_order
+
+    def test_multiple_wire_cuts_replaced(self):
+        """
+        Tests that all WireCut operators are replaced with MeasureNodes and
+        PrepareNodes with the correct order
+        """
+
+        wire_cut_1 = 0
+        wire_cut_2 = "a"
+        wire_cut_3 = 2
+        wire_cut_num = [wire_cut_1, wire_cut_2, wire_cut_3]
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.432, wires=0)
+            qml.RY(0.543, wires="a")
+            qml.WireCut(wires=wire_cut_1)
+            qml.CNOT(wires=[0, "a"])
+            qml.RZ(0.240, wires=0)
+            qml.RZ(0.133, wires="a")
+            qml.WireCut(wires=wire_cut_2)
+            qml.CNOT(wires=["a", 2])
+            qml.RX(0.432, wires="a")
+            qml.WireCut(wires=wire_cut_3)
+            qml.CNOT(wires=[2, 3])
+            qml.RY(0.543, wires=2)
+            qml.RZ(0.876, wires=3)
+            qml.expval(qml.PauliZ(wires=[0]))
+
+        g = qcut.tape_to_graph(tape)
+        node_data = list(g.nodes(data=True))
+
+        wire_cut_order = [order for op, order in node_data if op.name == "WireCut"]
+
+        qcut.replace_wire_cut_nodes(g)
+        new_node_data = list(g.nodes(data=True))
+        op_names = [op.name for op, order in new_node_data]
+
+        assert "WireCut" not in op_names
+        assert op_names.count("MeasureNode") == 3
+        assert op_names.count("PrepareNode") == 3
+
+        measure_counter = prepare_counter = 0
+
+        for op, order in new_node_data:
+            if op.name == "MeasureNode":
+                assert op.wires.tolist() == [wire_cut_num[measure_counter]]
+                assert order == wire_cut_order[measure_counter]
+                measure_counter += 1
+            elif op.name == "PrepareNode":
+                assert op.wires.tolist() == [wire_cut_num[prepare_counter]]
+                assert order == wire_cut_order[prepare_counter]
+                prepare_counter += 1
+
+    def test_successor_and_predecessor(self):
+        """
+        Tests the successor of the MeasureNode is the PrepareNode and the
+        predecessor of the PrepareNode is the MeasureNode
+        """
+        wire_cut_num = 1
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.432, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(0.133, wires=1)
+            qml.WireCut(wires=wire_cut_num)
+            qml.CNOT(wires=[1, 2])
+            qml.RY(0.543, wires=2)
+            qml.expval(qml.PauliZ(wires=[0]))
+
+        g = qcut.tape_to_graph(tape)
+        qcut.replace_wire_cut_nodes(g)
+
+        nodes = list(g.nodes)
+
+        for node in nodes:
+            if node.name == "MeasureNode":
+                succ = list(g.succ[node])[0]
+                pred = list(g.pred[node])[0]
+                assert succ.name == "PrepareNode"
+                assert pred.name == "RZ"
+            if node.name == "PrepareNode":
+                succ = list(g.succ[node])[0]
+                pred = list(g.pred[node])[0]
+                assert succ.name == "CNOT"
+                assert pred.name == "MeasureNode"
+
+    def test_wirecut_has_no_predecessor(self):
+        """
+        Tests a wirecut is replaced if it is the first operation in the tape
+        i.e if it has no predecessor
+        """
+
+        with qml.tape.QuantumTape() as tape:
+            qml.WireCut(wires=0)
+            qml.RX(0.432, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(0.133, wires=1)
+            qml.expval(qml.PauliZ(wires=[0]))
+
+        g = qcut.tape_to_graph(tape)
+        node_data = list(g.nodes(data=True))
+
+        wire_cut_order = {"order": 0}
+
+        qcut.replace_wire_cut_nodes(g)
+        new_node_data = list(g.nodes(data=True))
+        op_names = [op.name for op, order in new_node_data]
+
+        assert "WireCut" not in op_names
+        assert "MeasureNode" in op_names
+        assert "PrepareNode" in op_names
+
+        for op, order in new_node_data:
+            if op.name == "MeasureNode":
+                assert order == {"order": 0}
+                pred = list(g.pred[op])
+                assert pred == []
+            elif op.name == "PrepareNode":
+                assert order == {"order": 0}
+
+    def test_wirecut_has_no_successor(self):
+        """
+        Tests a wirecut is replaced if it is the last operation in the tape
+        i.e if it has no successor
+        """
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.432, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(0.133, wires=1)
+            qml.WireCut(wires=0)
+
+        g = qcut.tape_to_graph(tape)
+        node_data = list(g.nodes(data=True))
+
+        wire_cut_order = {"order": 3}
+
+        qcut.replace_wire_cut_nodes(g)
+        new_node_data = list(g.nodes(data=True))
+        op_names = [op.name for op, order in new_node_data]
+
+        assert "WireCut" not in op_names
+        assert "MeasureNode" in op_names
+        assert "PrepareNode" in op_names
+
+        for op, order in new_node_data:
+            if op.name == "MeasureNode":
+                assert order == {"order": 3}
+            elif op.name == "PrepareNode":
+                assert order == {"order": 3}
+                succ = list(g.succ[op])
+                assert succ == []
+
+    def test_multi_wire_wirecut_successor_and_predecessor(self):
+        """
+        Tests the successors and predecessors of a multi-wire wirecut are
+        correct
+        """
+        wire_cut_num = 1
+
+        with qml.tape.QuantumTape() as tape:
+            qml.CNOT(wires=[0, "a"])
+            qml.CNOT(wires=["a", 2])
+            qml.WireCut(wires=[0, "a", 2])
+            qml.CZ(wires=[0, 2])
+            qml.Toffoli(wires=[0, "a", 2])
+
+        g = qcut.tape_to_graph(tape)
+        qcut.replace_wire_cut_nodes(g)
+
+        nodes = list(g.nodes)
+        measure_nodes = [node for node in nodes if node.name == "MeasureNode"]
+        prepare_nodes = [node for node in nodes if node.name == "PrepareNode"]
+
+        assert len(measure_nodes) == len(prepare_nodes) == 3
+
+        expected_meas_pred_wires = [[0, "a"], ["a", 2], ["a", 2]]
+        expected_meas_pred_name = [
+            "CNOT",
+            "CNOT",
+        ] * len(measure_nodes)
+        expected_meas_succ_wires = [[0], ["a"], [2]]
+        expected_meas_succ_name = ["PrepareNode"] * len(measure_nodes)
+
+        expected_prep_pred_wires = expected_meas_succ_wires
+        exepeted_prep_pred_name = ["MeasureNode"] * len(measure_nodes)
+        expected_prep_succ_wires = [[0, 2], [0, "a", 2], [0, 2]]
+        expected_prep_succ_name = ["CZ", "Toffoli", "CZ"]
+
+        measure_pred = [list(g.pred[node])[0] for node in measure_nodes]
+        measure_succ = [list(g.succ[node])[0] for node in measure_nodes]
+        prep_pred = [list(g.pred[node])[0] for node in prepare_nodes]
+        prep_succ = [list(g.succ[node])[0] for node in prepare_nodes]
+
+        assert len(measure_nodes) == len(prepare_nodes) == 3
+
+        def compare_nodes(nodes, expected_wires, expected_names):
+
+            for node, exp_wire in zip(nodes, expected_wires):
+                assert node.wires.tolist() == exp_wire
+
+            for node, exp_name in zip(nodes, expected_names):
+                assert node.name == exp_name
+
+        compare_nodes(measure_pred, expected_meas_pred_wires, expected_meas_pred_name)
+        compare_nodes(measure_succ, expected_meas_succ_wires, expected_meas_succ_name)
+        compare_nodes(prep_pred, expected_prep_pred_wires, exepeted_prep_pred_name)
+        compare_nodes(prep_succ, expected_prep_succ_wires, expected_prep_succ_name)
+
+        for node in measure_nodes + prepare_nodes:
+            in_edges = list(g.in_edges(node, data="wire"))
+            out_edges = list(g.out_edges(node, data="wire"))
+            assert len(in_edges) == 1
+            assert len(out_edges) == 1
+
+            _, _, wire_label_in = in_edges[0]
+            _, _, wire_label_out = out_edges[0]
+
+            assert wire_label_in == wire_label_out == node.wires.tolist()[0]