Repository: PennyLaneAI/pennylane
Pull request #2153 (2022-02-02 16:37:17)

Cut graphs can be fragmented into subgraphs and communication graph
**Context:**
The circuit cutting pipeline consists of manually placing a [wirecut operator](https://github.com/PennyLaneAI/pennylane/pull/2093) during construction of a circuit, converting this circuit to a [graph](https://github.com/PennyLaneAI/pennylane/pull/2107) and then [replacing](https://github.com/PennyLaneAI/pennylane/pull/2124) the wirecut nodes with measure and prepare nodes. At this stage the graph is ready to be fragmented into subgraphs. A communication graph is also returned, this is a graph where each node represents a fragment and edges denote the flow of qubits between fragments. This is required for reconstruction in following steps. 

**Description of the Change:**
A method for fragmenting cut graphs has been added

**Benefits:**
Facilitates the circuit cutting compiler


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2153](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (755663e) into [sc-13391-wirecut-nodes-can-be-replaced-with-measure](https://codecov.io/gh/PennyLaneAI/pennylane/commit/76befee1a373713a667539053e3a563cfab67711?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (76befee) will **decrease** coverage by `2.80%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/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/2153?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@                                   Coverage Diff                                   @@
##           sc-13391-wirecut-nodes-can-be-replaced-with-measure    #2153      +/-   ##
=======================================================================================
- Coverage                                                99.20%   96.40%   -2.81%     
=======================================================================================
  Files                                                      231      231              
  Lines                                                    17785    17806      +21     
=======================================================================================
- Hits                                                     17644    17166     -478     
- Misses                                                     141      640     +499     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153?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/2153/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%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4X2ppdC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `84.41% <0.00%> (-15.59%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.73% <0.00%> (-12.27%)` | :arrow_down: |
| [pennylane/math/is\_independent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvaXNfaW5kZXBlbmRlbnQucHk=) | `88.50% <0.00%> (-11.50%)` | :arrow_down: |
| ... and [14 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153/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/2153?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/2153?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [76befee...755663e](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2153?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).

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6c2ef3881..3d38c5d02 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -14,6 +14,10 @@
   and `PrepareNode` placeholders has been added.
   [(#2124)](https://github.com/PennyLaneAI/pennylane/pull/2124)
   
+  A method has been added that takes a directed multigraph with `MeasureNode` and
+  `PrepareNode` placeholders and fragments into subgraphs and a communication graph.
+  [(#2153)](https://github.com/PennyLaneAI/pennylane/pull/2153)
+
 <h3>Improvements</h3>
 
 <h3>Breaking changes</h3>
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 0f57ea8f6..23ae30c6b 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -82,6 +82,19 @@ both transforms, and decompositions within the larger PennyLane codebase.
     ~transforms.two_qubit_decomposition
     ~transforms.set_decomposition
 
+Transforms for circuit cutting
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The following are utility functions that compose the circuit cutting transform.
+
+.. autosummary::
+    :toctree: api
+
+    ~transforms.tape_to_graph
+    ~transforms.replace_wire_cut_node
+    ~transforms.replace_wire_cut_nodes
+    ~transforms.fragment_graph
+
 Transforms that act on tapes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -153,3 +166,4 @@ from .tape_expand import (
     create_decomp_expand_fn,
     set_decomposition,
 )
+from .qcut import tape_to_graph, replace_wire_cut_node, replace_wire_cut_nodes, fragment_graph
diff --git a/pennylane/transforms/qcut.py b/pennylane/transforms/qcut.py
index 2f05df794..8e1b059ac 100644
--- a/pennylane/transforms/qcut.py
+++ b/pennylane/transforms/qcut.py
@@ -15,8 +15,9 @@
 This module provides the circuit cutting functionality that allows large
 circuits to be distributed across multiple devices.
 """
+from typing import Tuple
 
-from networkx import MultiDiGraph
+from networkx import MultiDiGraph, weakly_connected_components
 from pennylane.measure import MeasurementProcess
 from pennylane.operation import Operation, Operator, Tensor
 from pennylane.ops.qubit.non_parametric_ops import WireCut
@@ -53,8 +54,6 @@ def replace_wire_cut_node(node: WireCut, graph: MultiDiGraph):
 
     .. code-block:: python
 
-        from pennylane.transforms import qcut
-
         wire_cut = qml.WireCut(wires=0)
 
         with qml.tape.QuantumTape() as tape:
@@ -65,8 +64,8 @@ def replace_wire_cut_node(node: WireCut, graph: MultiDiGraph):
 
     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)
+    >>> graph = qml.transforms.tape_to_graph(tape)
+    >>> qml.transforms.replace_wire_cut_node(wire_cut, graph)
     """
     predecessors = graph.pred[node]
     successors = graph.succ[node]
@@ -122,8 +121,6 @@ def replace_wire_cut_nodes(graph: MultiDiGraph):
 
     .. 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])
@@ -140,9 +137,8 @@ def replace_wire_cut_nodes(graph: MultiDiGraph):
 
     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)
-
+    >>> graph = qml.transforms.tape_to_graph(tape)
+    >>> qml.transforms.replace_wire_cut_nodes(graph)
     """
     for op in list(graph.nodes):
         if isinstance(op, WireCut):
@@ -169,8 +165,8 @@ def tape_to_graph(tape: QuantumTape) -> MultiDiGraph:
         tape (QuantumTape): tape to be converted into a directed multigraph
 
     Returns:
-        graph (MultiDiGraph): a directed multigraph that captures the circuit
-        structure of the input tape
+        MultiDiGraph: a directed multigraph that captures the circuit structure
+        of the input tape
 
     **Example**
 
@@ -186,7 +182,7 @@ def tape_to_graph(tape: QuantumTape) -> MultiDiGraph:
 
     Its corresponding circuit graph can be found using
 
-    >>> tape_to_graph(tape)
+    >>> qml.transforms.tape_to_graph(tape)
     <networkx.classes.multidigraph.MultiDiGraph at 0x7fe41cbd7210>
     """
     graph = MultiDiGraph()
@@ -210,3 +206,79 @@ def tape_to_graph(tape: QuantumTape) -> MultiDiGraph:
             order += 1
 
     return graph
+
+
+def fragment_graph(graph: MultiDiGraph) -> Tuple[Tuple[MultiDiGraph], MultiDiGraph]:
+    """
+    Fragments a graph into a collection of subgraphs as well as returning
+    the communication/`quotient <https://en.wikipedia.org/wiki/Quotient_graph>`__
+    graph. Each node of the communication graph represents a fragment and the edges
+    denote the flow of qubits between fragments.
+
+    Args:
+        graph (MultiDiGraph): directed multigraph containing measure and prepare
+            nodes at cut locations
+
+    Returns:
+        Tuple[Tuple[MultiDiGraph], MultiDiGraph]: the subgraphs of the cut graph
+        and the communication graph.
+
+    **Example**
+
+    Consider the following circuit with manually-placed wire cuts:
+
+    .. code-block:: python
+
+        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 corresponding graph, remove all the wire cut nodes, and
+    find the subgraphs and communication graph by using:
+
+    >>> graph = qml.transforms.tape_to_graph(tape)
+    >>> qml.transforms.replace_wire_cut_nodes(graph)
+    >>> qml.transforms.fragment_graph(graph)
+    ((<networkx.classes.multidigraph.MultiDiGraph object at 0x7fb3b2311940>,
+      <networkx.classes.multidigraph.MultiDiGraph object at 0x7fb3b2311c10>,
+      <networkx.classes.multidigraph.MultiDiGraph object at 0x7fb3b23e2820>,
+      <networkx.classes.multidigraph.MultiDiGraph object at 0x7fb3b23e27f0>),
+     <networkx.classes.multidigraph.MultiDiGraph object at 0x7fb3b23e26a0>)
+    """
+
+    graph_copy = graph.copy()
+
+    cut_edges = []
+
+    for node1, node2, wire in graph.edges:
+        if isinstance(node1, MeasureNode):
+            assert isinstance(node2, PrepareNode)
+            cut_edges.append((node1, node2, wire))
+            graph_copy.remove_edge(node1, node2, key=wire)
+
+    subgraph_nodes = weakly_connected_components(graph_copy)
+    subgraphs = tuple(graph_copy.subgraph(n) for n in subgraph_nodes)
+
+    communication_graph = MultiDiGraph()
+    communication_graph.add_nodes_from(range(len(subgraphs)))
+
+    for node1, node2, wire in cut_edges:
+        for i, subgraph in enumerate(subgraphs):
+            if subgraph.has_node(node1):
+                start_fragment = i
+            if subgraph.has_node(node2):
+                end_fragment = i
+
+        communication_graph.add_edge(start_fragment, end_fragment, pair=(node1, node2, wire))
+
+    return subgraphs, communication_graph


Test patch:
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index 4dfac069b..35cfff3af 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -1,5 +1,5 @@
 # Copyright 2022 Xanadu Quantum Technologies Inc.
-#
+
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
@@ -28,6 +28,52 @@ with qml.tape.QuantumTape() as tape:
     qml.RZ(0.133, wires="a")
     qml.expval(qml.PauliZ(wires=[0]))
 
+with qml.tape.QuantumTape() as multi_cut_tape:
+    qml.RX(0.432, wires=0)
+    qml.RY(0.543, wires="a")
+    qml.WireCut(wires=0)
+    qml.CNOT(wires=[0, "a"])
+    qml.RZ(0.240, wires=0)
+    qml.RZ(0.133, wires="a")
+    qml.WireCut(wires="a")
+    qml.CNOT(wires=["a", 2])
+    qml.RX(0.432, wires="a")
+    qml.WireCut(wires=2)
+    qml.CNOT(wires=[2, 3])
+    qml.RY(0.543, wires=2)
+    qml.RZ(0.876, wires=3)
+    qml.expval(qml.PauliZ(wires=[0]))
+
+
+def compare_nodes(nodes, expected_wires, expected_names):
+    """Helper function to compare nodes of directed multigraph"""
+
+    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
+
+
+def compare_fragment_nodes(node_data, expected_data):
+    """Helper function to compare nodes of fragment graphs"""
+
+    expected = [(exp_data[0].name, exp_data[0].wires, exp_data[1]) for exp_data in expected_data]
+
+    for data in node_data:
+        # The exact ordering of node_data varies on each call
+        assert (data[0].name, data[0].wires, data[1]) in expected
+
+
+def compare_fragment_edges(edge_data, expected_data):
+    """Helper function to compare fragment edges"""
+
+    expected = [(exp_data[0].name, exp_data[1].name, exp_data[2]) for exp_data in expected_data]
+
+    for data in edge_data:
+        # The exact ordering of edge_data varies on each call
+        assert (data[0].name, data[1].name, data[2]) in expected
+
 
 class TestTapeToGraph:
     """
@@ -477,14 +523,6 @@ class TestReplaceWireCut:
 
         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)
@@ -500,3 +538,195 @@ class TestReplaceWireCut:
             _, _, wire_label_out = out_edges[0]
 
             assert wire_label_in == wire_label_out == node.wires.tolist()[0]
+
+
+class TestFragmentGraph:
+    """
+    Tests that a cut graph is fragmented into subgraphs correctly
+    """
+
+    def test_subgraphs_of_multi_wirecut(self):
+        """
+        Tests that the subgraphs of a graph with multiple wirecuts contain the
+        correct nodes and edges
+        """
+
+        g = qcut.tape_to_graph(multi_cut_tape)
+        qcut.replace_wire_cut_nodes(g)
+        subgraphs, communication_graph = qcut.fragment_graph(g)
+
+        assert len(subgraphs) == 4
+
+        sub_0_expected_nodes = [
+            (qcut.MeasureNode(wires=[0]), {"order": 2}),
+            (qml.RX(0.432, wires=[0]), {"order": 0}),
+        ]
+        sub_1_expected_nodes = [
+            (qml.RY(0.543, wires=["a"]), {"order": 1}),
+            (qcut.PrepareNode(wires=[0]), {"order": 2}),
+            (qcut.MeasureNode(wires=["a"]), {"order": 6}),
+            (qml.RZ(0.24, wires=[0]), {"order": 4}),
+            (qml.CNOT(wires=[0, "a"]), {"order": 3}),
+            (qml.expval(qml.PauliZ(wires=[0])), {"order": 13}),
+            (qml.RZ(0.133, wires=["a"]), {"order": 5}),
+        ]
+        sub_2_expected_nodes = [
+            (qml.RX(0.432, wires=["a"]), {"order": 8}),
+            (qcut.MeasureNode(wires=[2]), {"order": 9}),
+            (qcut.PrepareNode(wires=["a"]), {"order": 6}),
+            (qml.CNOT(wires=["a", 2]), {"order": 7}),
+        ]
+        sub_3_expected_nodes = [
+            (qml.CNOT(wires=[2, 3]), {"order": 10}),
+            (qml.RY(0.543, wires=[2]), {"order": 11}),
+            (qcut.PrepareNode(wires=[2]), {"order": 9}),
+            (qml.RZ(0.876, wires=[3]), {"order": 12}),
+        ]
+        expected_nodes = [
+            sub_0_expected_nodes,
+            sub_1_expected_nodes,
+            sub_2_expected_nodes,
+            sub_3_expected_nodes,
+        ]
+
+        sub_0_expected_edges = [
+            (qml.RX(0.432, wires=[0]), qcut.MeasureNode(wires=[0]), {"wire": 0})
+        ]
+        sub_1_expected_edges = [
+            (qcut.PrepareNode(wires=[0]), qml.CNOT(wires=[0, "a"]), {"wire": 0}),
+            (qml.RZ(0.24, wires=[0]), qml.expval(qml.PauliZ(wires=[0])), {"wire": 0}),
+            (qml.RZ(0.133, wires=["a"]), qcut.MeasureNode(wires=["a"]), {"wire": "a"}),
+            (qml.CNOT(wires=[0, "a"]), qml.RZ(0.24, wires=[0]), {"wire": 0}),
+            (qml.CNOT(wires=[0, "a"]), qml.RZ(0.133, wires=["a"]), {"wire": "a"}),
+            (qml.RY(0.543, wires=["a"]), qml.CNOT(wires=[0, "a"]), {"wire": "a"}),
+        ]
+        sub_2_expected_edges = [
+            (qcut.PrepareNode(wires=["a"]), qml.CNOT(wires=["a", 2]), {"wire": "a"}),
+            (qml.CNOT(wires=["a", 2]), qml.RX(0.432, wires=["a"]), {"wire": "a"}),
+            (qml.CNOT(wires=["a", 2]), qcut.MeasureNode(wires=[2]), {"wire": 2}),
+        ]
+        sub_3_expected_edges = [
+            (qcut.PrepareNode(wires=[2]), qml.CNOT(wires=[2, 3]), {"wire": 2}),
+            (qml.CNOT(wires=[2, 3]), qml.RY(0.543, wires=[2]), {"wire": 2}),
+            (qml.CNOT(wires=[2, 3]), qml.RZ(0.876, wires=[3]), {"wire": 3}),
+        ]
+        expected_edges = [
+            sub_0_expected_edges,
+            sub_1_expected_edges,
+            sub_2_expected_edges,
+            sub_3_expected_edges,
+        ]
+
+        for subgraph, expected_n in zip(subgraphs, expected_nodes):
+            compare_fragment_nodes(list(subgraph.nodes(data=True)), expected_n)
+
+        for subgraph, expected_e in zip(subgraphs, expected_edges):
+            compare_fragment_edges(list(subgraph.edges(data=True)), expected_e)
+
+    def test_communication_graph(self):
+        """
+        Tests that the communication graph contains the correct nodes and edges
+        """
+
+        g = qcut.tape_to_graph(multi_cut_tape)
+        qcut.replace_wire_cut_nodes(g)
+        subgraphs, communication_graph = qcut.fragment_graph(g)
+
+        assert list(communication_graph.nodes) == list(range(4))
+
+        expected_edge_data = [
+            (0, 1, {"pair": (qcut.MeasureNode(wires=[0]), qcut.PrepareNode(wires=[0]))}),
+            (1, 2, {"pair": (qcut.MeasureNode(wires=["a"]), qcut.PrepareNode(wires=["a"]))}),
+            (2, 3, {"pair": (qcut.MeasureNode(wires=[2]), qcut.PrepareNode(wires=[2]))}),
+        ]
+        edge_data = list(communication_graph.edges(data=True))
+
+        for edge, exp_edge in zip(edge_data, expected_edge_data):
+            assert edge[0] == exp_edge[0]
+            assert edge[1] == exp_edge[1]
+
+            for node, exp_node in zip(edge[2]["pair"], exp_edge[2]["pair"]):
+                assert node.name == exp_node.name
+                assert node.wires.tolist() == exp_node.wires.tolist()
+
+    def test_fragment_wirecut_first_and_last(self):
+        """
+        Tests a circuit with wirecut at the start and end is fragmented
+        correctly
+        """
+
+        with qml.tape.QuantumTape() as tape:
+            qml.WireCut(wires=0)
+            qml.RX(0.432, wires=0)
+            qml.RY(0.543, wires="a")
+            qml.WireCut(wires="a")
+
+        g = qcut.tape_to_graph(tape)
+        qcut.replace_wire_cut_nodes(g)
+        subgraphs, communication_graph = qcut.fragment_graph(g)
+
+        sub_0_expected_nodes = [
+            (qcut.PrepareNode(wires=[0]), {"order": 0}),
+            (qml.RX(0.432, wires=[0]), {"order": 1}),
+        ]
+        sub_1_expected_nodes = [
+            (qcut.MeasureNode(wires=["a"]), {"order": 3}),
+            (qml.RY(0.543, wires=["a"]), {"order": 2}),
+        ]
+        sub_2_expected_nodes = [(qcut.MeasureNode(wires=[0]), {"order": 0})]
+        sub_3_expected_nodes = [(qcut.PrepareNode(wires=["a"]), {"order": 3})]
+
+        expected_nodes = [
+            sub_0_expected_nodes,
+            sub_1_expected_nodes,
+            sub_2_expected_nodes,
+            sub_3_expected_nodes,
+        ]
+
+        sub_0_expected_edges = [
+            (qcut.PrepareNode(wires=[0]), qml.RX(0.432, wires=[0]), {"wire": 0})
+        ]
+        sub_1_expected_edges = [
+            (qml.RY(0.543, wires=["a"]), qcut.MeasureNode(wires=["a"]), {"wire": "a"})
+        ]
+        sub_2_expected_edges = []
+        sub_3_expected_edges = []
+
+        expected_edges = [
+            sub_0_expected_edges,
+            sub_1_expected_edges,
+            sub_2_expected_edges,
+            sub_3_expected_edges,
+        ]
+
+        for subgraph, expected_n in zip(subgraphs, expected_nodes):
+            compare_fragment_nodes(list(subgraph.nodes(data=True)), expected_n)
+
+        for subgraph, expected_e in zip(subgraphs, expected_edges):
+            compare_fragment_edges(list(subgraph.edges(data=True)), expected_e)
+
+    def test_communication_graph_persistence(self):
+        """
+        Tests that when `fragment_graph` is repeatedly applied the
+        communication graph is the same each time.
+        """
+
+        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=1)
+            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)
+        qcut.replace_wire_cut_nodes(g)
+        subgraphs_0, communication_graph_0 = qcut.fragment_graph(g)
+        subgraphs_1, communication_graph_1 = qcut.fragment_graph(g)
+
+        assert communication_graph_0.nodes == communication_graph_1.nodes
+        assert communication_graph_0.edges == communication_graph_1.edges