Repository: PennyLaneAI/pennylane
Pull request #2158 (2022-02-02 22:13:38)

Circuit cutting: Add differentiable tensor contraction
**Context:**

Provides a differentiable tensor contraction function `contract_tensors` to be used as part of the circuit cutting workflow.

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.
trbromley: [sc-13397]
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2158?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2158](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2158?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (31ffc38) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/1e480163db3d27e983dca079a1324476bf4e9c32?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1e48016) will **decrease** coverage by `3.00%`.
> The diff coverage is `5.26%`.

> :exclamation: Current head 31ffc38 differs from pull request most recent head 14fcb1f. Consider uploading reports for the commit 14fcb1f to get more accurate results

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

```diff
@@            Coverage Diff             @@
##           master    #2158      +/-   ##
==========================================
- Coverage   99.20%   96.20%   -3.01%     
==========================================
  Files         231      231              
  Lines       17785    17823      +38     
==========================================
- Hits        17644    17147     -497     
- Misses        141      676     +535     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2158?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/2158/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvcWN1dC5weQ==) | `65.04% <5.26%> (-34.96%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2158/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/2158/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/2158/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/2158/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/2158/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/2158/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/2158/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/2158/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/2158/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/2158/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/2158?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/2158?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [1e48016...14fcb1f](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2158?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).

trbromley: Thanks @josh146 for the review!

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 299d7e265..444b97d8d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -18,6 +18,10 @@
   `PrepareNode` placeholders and fragments into subgraphs and a communication graph.
   [(#2153)](https://github.com/PennyLaneAI/pennylane/pull/2153)
 
+  A differentiable tensor contraction function `contract_tensors` has been
+  added.
+  [(#2158)](https://github.com/PennyLaneAI/pennylane/pull/2158)
+
 <h3>Improvements</h3>
 
 <h3>Breaking changes</h3>
@@ -40,4 +44,4 @@
 
 This release contains contributions from (in alphabetical order):
 
-Anthony Hayes, Josh Izaac, David Wierichs
+Thomas Bromley, Anthony Hayes, Josh Izaac, David Wierichs
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 23ae30c6b..5abfcf05a 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -94,6 +94,7 @@ The following are utility functions that compose the circuit cutting transform.
     ~transforms.replace_wire_cut_node
     ~transforms.replace_wire_cut_nodes
     ~transforms.fragment_graph
+    ~transforms.contract_tensors
 
 Transforms that act on tapes
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -166,4 +167,10 @@ 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
+from .qcut import (
+    tape_to_graph,
+    replace_wire_cut_node,
+    replace_wire_cut_nodes,
+    fragment_graph,
+    contract_tensors,
+)
diff --git a/pennylane/transforms/qcut.py b/pennylane/transforms/qcut.py
index 8e1b059ac..762caac79 100644
--- a/pennylane/transforms/qcut.py
+++ b/pennylane/transforms/qcut.py
@@ -15,9 +15,12 @@
 This module provides the circuit cutting functionality that allows large
 circuits to be distributed across multiple devices.
 """
-from typing import Tuple
+import string
+from typing import Sequence, Tuple
 
 from networkx import MultiDiGraph, weakly_connected_components
+
+import pennylane as qml
 from pennylane.measure import MeasurementProcess
 from pennylane.operation import Operation, Operator, Tensor
 from pennylane.ops.qubit.non_parametric_ops import WireCut
@@ -282,3 +285,146 @@ def fragment_graph(graph: MultiDiGraph) -> Tuple[Tuple[MultiDiGraph], MultiDiGra
         communication_graph.add_edge(start_fragment, end_fragment, pair=(node1, node2, wire))
 
     return subgraphs, communication_graph
+
+
+def _get_symbol(i):
+    """Finds the i-th ASCII symbol. Works for lowercase and uppercase letters, allowing i up to
+    51."""
+    if i >= len(string.ascii_letters):
+        raise ValueError(
+            "Set the use_opt_einsum argument to True when applying more than "
+            f"{len(string.ascii_letters)} wire cuts to a circuit"
+        )
+    return string.ascii_letters[i]
+
+
+# pylint: disable=too-many-branches
+def contract_tensors(
+    tensors: Sequence,
+    communication_graph: MultiDiGraph,
+    prepare_nodes: Sequence[Sequence[PrepareNode]],
+    measure_nodes: Sequence[Sequence[MeasureNode]],
+    use_opt_einsum: bool = False,
+):
+    r"""Contract tensors according to the edges specified in the communication graph.
+
+    .. note::
+
+        This function is designed for use as part of the circuit cutting workflow. Check out the
+        :doc:`transforms </code/qml_transforms>` page for more details.
+
+    Consider the three tensors :math:`T^{(1)}`, :math:`T^{(2)}`, and :math:`T^{(3)}`, along with
+    their contraction equation
+
+    .. math::
+
+        \sum_{ijklmn} T^{(1)}_{ij,km} T^{(2)}_{kl,in} T^{(3)}_{mn,jl}
+
+    Each tensor is the result of the tomography of a circuit fragment and has some indices
+    corresponding to state preparations (marked by the indices before the comma) and some indices
+    corresponding to measurements (marked by the indices after the comma).
+
+    An equivalent representation of the contraction equation is to use a directed multigraph known
+    as the communication/quotient graph. In the communication graph, each tensor is assigned a node
+    and edges are added between nodes to mark a contraction along an index. The communication graph
+    resulting from the above contraction equation is a complete directed graph.
+
+    In the communication graph provided by :func:`fragment_graph`, edges are composed of
+    :class:`PrepareNode` and :class:`MeasureNode` pairs. To correctly map back to the contraction
+    equation, we must keep track of the order of preparation and measurement indices in each tensor.
+    This order is specified in the ``prepare_nodes`` and ``measure_nodes`` arguments.
+
+    Args:
+        tensors (Sequence): the tensors to be contracted
+        communication_graph (MultiDiGraph): the communication graph determining connectivity between
+            the tensors
+        prepare_nodes (Sequence[Sequence[PrepareNode]]): a sequence of size
+            ``len(communication_graph.nodes)`` that determines the order of preparation indices in
+            each tensor
+        measure_nodes (Sequence[Sequence[MeasureNode]]): a sequence of size
+            ``len(communication_graph.nodes)`` that determines the order of measurement indices in
+            each tensor
+        use_opt_einsum (bool): Determines whether to use the
+            `opt_einsum <https://dgasmith.github.io/opt_einsum/>`__ package. This package is useful
+            for tensor contractions of large networks but must be installed separately using, e.g.,
+            ``pip install opt_einsum``. Both settings for ``use_opt_einsum`` result in a
+            differentiable contraction.
+
+    Returns:
+        float or array-like: the result of contracting the tensor network
+
+    **Example**
+
+    We first set up the tensors and their corresponding :class:`~.PrepareNode` and
+    :class:`~.MeasureNode` orderings:
+
+    .. code-block:: python
+
+        from pennylane.transforms import qcut
+        import networkx as nx
+        import numpy as np
+
+        tensors = [np.arange(4), np.arange(4, 8)]
+        prep = [[], [qcut.PrepareNode(wires=0)]]
+        meas = [[qcut.MeasureNode(wires=0)], []]
+
+    The communication graph describing edges in the tensor network must also be constructed:
+
+    .. code-block:: python
+
+        graph = nx.MultiDiGraph([(0, 1, {"pair": (meas[0][0], prep[1][0])})])
+
+    The network can then be contracted using:
+
+    >>> qml.transforms.contract_tensors(tensors, graph, prep, meas)
+    38
+    """
+    # pylint: disable=import-outside-toplevel
+    if use_opt_einsum:
+        try:
+            from opt_einsum import contract, get_symbol
+        except ImportError as e:
+            raise ImportError(
+                "The opt_einsum package is required when use_opt_einsum is set to "
+                "True in the contract_tensors function. This package can be "
+                "installed using:\npip install opt_einsum"
+            ) from e
+    else:
+        contract = qml.math.einsum
+        get_symbol = _get_symbol
+
+    ctr = 0
+    tensor_indxs = [""] * len(communication_graph.nodes)
+
+    meas_map = {}
+
+    for i, (node, prep) in enumerate(zip(communication_graph.nodes, prepare_nodes)):
+        predecessors = communication_graph.pred[node]
+
+        for p in prep:
+            for _, pred_edges in predecessors.items():
+                for pred_edge in pred_edges.values():
+                    meas_op, prep_op = pred_edge["pair"]
+
+                    if p is prep_op:
+                        symb = get_symbol(ctr)
+                        ctr += 1
+                        tensor_indxs[i] += symb
+                        meas_map[meas_op] = symb
+
+    for i, (node, meas) in enumerate(zip(communication_graph.nodes, measure_nodes)):
+        successors = communication_graph.succ[node]
+
+        for m in meas:
+            for _, succ_edges in successors.items():
+                for succ_edge in succ_edges.values():
+                    meas_op, _ = succ_edge["pair"]
+
+                    if m is meas_op:
+                        symb = meas_map[meas_op]
+                        tensor_indxs[i] += symb
+
+    eqn = ",".join(tensor_indxs)
+    kwargs = {} if use_opt_einsum else {"like": tensors[0]}
+
+    return contract(eqn, *tensors, **kwargs)
diff --git a/requirements-ci.txt b/requirements-ci.txt
index 07bec466c..279c11a09 100644
--- a/requirements-ci.txt
+++ b/requirements-ci.txt
@@ -12,3 +12,4 @@ semantic_version==2.6
 dask[delayed]
 autoray
 matplotlib
+opt_einsum
diff --git a/requirements.txt b/requirements.txt
index dfeef6c57..8f17ae9da 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,4 +13,5 @@ semantic_version==2.6
 dask[delayed]==2021.10
 autoray>=0.2.5
 matplotlib==3.4
-black>=21
\ No newline at end of file
+black>=21
+opt_einsum==3.3
\ No newline at end of file


Test patch:
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index 35cfff3af..1ee529f7d 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -14,9 +14,14 @@
 """
 Unit tests for the `pennylane.qcut` package.
 """
-import numpy as np
-import pennylane as qml
+import string
+import sys
+
 import pytest
+from networkx import MultiDiGraph
+
+import pennylane as qml
+from pennylane import numpy as np
 from pennylane.transforms import qcut
 
 with qml.tape.QuantumTape() as tape:
@@ -730,3 +735,214 @@ class TestFragmentGraph:
 
         assert communication_graph_0.nodes == communication_graph_1.nodes
         assert communication_graph_0.edges == communication_graph_1.edges
+
+
+class TestContractTensors:
+    """Tests for the contract_tensors function"""
+
+    t = [np.arange(4), np.arange(4, 8)]
+    m = [[qcut.MeasureNode(wires=0)], []]
+    p = [[], [qcut.PrepareNode(wires=0)]]
+    edge_dict = {"pair": (m[0][0], p[1][0])}
+    g = MultiDiGraph([(0, 1, edge_dict)])
+    expected_result = np.dot(*t)
+
+    @pytest.mark.parametrize("use_opt_einsum", [False, True])
+    def test_basic(self, use_opt_einsum):
+        """Test that the correct answer is returned for a simple contraction scenario"""
+        if use_opt_einsum:
+            pytest.importorskip("opt_einsum")
+        res = qcut.contract_tensors(self.t, self.g, self.p, self.m, use_opt_einsum=use_opt_einsum)
+
+        assert np.allclose(res, self.expected_result)
+
+    def test_fail_import(self, monkeypatch):
+        """Test if an ImportError is raised when opt_einsum is requested but not installed"""
+
+        with monkeypatch.context() as m:
+            m.setitem(sys.modules, "opt_einsum", None)
+
+            with pytest.raises(ImportError, match="The opt_einsum package is required"):
+                qcut.contract_tensors(self.t, self.g, self.p, self.m, use_opt_einsum=True)
+
+    def test_run_out_of_symbols(self, monkeypatch):
+        """Test if a ValueError is raised when there are not enough symbols in the
+        use_opt_einsum = False setting"""
+
+        with monkeypatch.context() as m:
+            m.setattr(string, "ascii_letters", "")
+            with pytest.raises(ValueError, match="Set the use_opt_einsum argument to True"):
+                qcut.contract_tensors(self.t, self.g, self.p, self.m, use_opt_einsum=False)
+
+    params = [0.3, 0.5]
+
+    expected_grad_0 = (
+        np.cos(params[0]) * np.cos(params[1])
+        + 2 * np.cos(params[0]) * np.sin(params[0]) * np.cos(params[1]) ** 2
+        + 3 * np.cos(params[0]) * np.sin(params[0]) ** 2 * np.cos(params[1]) ** 3
+    )
+    expected_grad_1 = (
+        -np.sin(params[0]) * np.sin(params[1])
+        - 2 * np.sin(params[0]) ** 2 * np.sin(params[1]) * np.cos(params[1])
+        - 3 * np.sin(params[0]) ** 3 * np.sin(params[1]) * np.cos(params[1]) ** 2
+    )
+    expected_grad = np.array([expected_grad_0, expected_grad_1])
+
+    @pytest.mark.parametrize("use_opt_einsum", [True, False])
+    def test_basic_grad_autograd(self, use_opt_einsum):
+        """Test if the basic contraction is differentiable using the autograd interface"""
+        if use_opt_einsum:
+            pytest.importorskip("opt_einsum")
+
+        def contract(params):
+            t1 = np.asarray([np.sin(params[0]) ** i for i in range(4)])
+            t2 = np.asarray([np.cos(params[1]) ** i for i in range(4)])
+            t = [t1, t2]
+            r = qcut.contract_tensors(t, self.g, self.p, self.m, use_opt_einsum=use_opt_einsum)
+            return r
+
+        params = np.array(self.params, requires_grad=True)
+        grad = qml.grad(contract)(params)
+
+        assert np.allclose(grad, self.expected_grad)
+
+    @pytest.mark.usefixtures("skip_if_no_torch_support")
+    @pytest.mark.parametrize("use_opt_einsum", [True, False])
+    def test_basic_grad_torch(self, use_opt_einsum):
+        """Test if the basic contraction is differentiable using the torch interface"""
+        if use_opt_einsum:
+            pytest.importorskip("opt_einsum")
+        import torch
+
+        params = torch.tensor(self.params, requires_grad=True)
+
+        t1 = torch.stack([torch.sin(params[0]) ** i for i in range(4)])
+        t2 = torch.stack([torch.cos(params[1]) ** i for i in range(4)])
+        t = [t1, t2]
+        r = qcut.contract_tensors(t, self.g, self.p, self.m, use_opt_einsum=use_opt_einsum)
+
+        r.backward()
+        grad = params.grad
+
+        assert np.allclose(grad, self.expected_grad)
+
+    @pytest.mark.usefixtures("skip_if_no_tf_support")
+    @pytest.mark.parametrize("use_opt_einsum", [True, False])
+    def test_basic_grad_tf(self, use_opt_einsum):
+        """Test if the basic contraction is differentiable using the tf interface"""
+        if use_opt_einsum:
+            pytest.importorskip("opt_einsum")
+        import tensorflow as tf
+
+        params = tf.Variable(self.params)
+
+        with tf.GradientTape() as tape:
+            t1 = tf.stack([tf.sin(params[0]) ** i for i in range(4)])
+            t2 = tf.stack([tf.cos(params[1]) ** i for i in range(4)])
+            t = [t1, t2]
+            r = qcut.contract_tensors(t, self.g, self.p, self.m, use_opt_einsum=use_opt_einsum)
+
+        grad = tape.gradient(r, params)
+
+        assert np.allclose(grad, self.expected_grad)
+
+    @pytest.mark.usefixtures("skip_if_no_jax_support")
+    @pytest.mark.parametrize("use_opt_einsum", [True, False])
+    def test_basic_grad_jax(self, use_opt_einsum):
+        """Test if the basic contraction is differentiable using the jax interface"""
+        if use_opt_einsum:
+            pytest.importorskip("opt_einsum")
+        import jax
+        from jax import numpy as np
+
+        params = np.array(self.params)
+
+        def contract(params):
+            t1 = np.stack([np.sin(params[0]) ** i for i in range(4)])
+            t2 = np.stack([np.cos(params[1]) ** i for i in range(4)])
+            t = [t1, t2]
+            r = qcut.contract_tensors(t, self.g, self.p, self.m, use_opt_einsum=use_opt_einsum)
+            return r
+
+        grad = jax.grad(contract)(params)
+
+        assert np.allclose(grad, self.expected_grad)
+
+    @pytest.mark.parametrize("use_opt_einsum", [True, False])
+    def test_advanced(self, mocker, use_opt_einsum):
+        """Test if the contraction works as expected for a more complicated example based
+        upon the circuit:
+
+        with qml.tape.QuantumTape() as tape:
+            qml.QubitUnitary(np.eye(2 ** 3), wires=[0, 1, 2])
+            qml.QubitUnitary(np.eye(2 ** 2), wires=[3, 4])
+
+            qml.Barrier(wires=0)
+            qml.WireCut(wires=[1, 2, 3])
+
+            qml.QubitUnitary(np.eye(2 ** 1), wires=[0])
+            qml.QubitUnitary(np.eye(2 ** 4), wires=[1, 2, 3, 4])
+
+            qml.WireCut(wires=[1, 2, 3, 4])
+
+            qml.QubitUnitary(np.eye(2 ** 3), wires=[0, 1, 2])
+            qml.QubitUnitary(np.eye(2 ** 2), wires=[3, 4])
+        """
+        if use_opt_einsum:
+            opt_einsum = pytest.importorskip("opt_einsum")
+            spy = mocker.spy(opt_einsum, "contract")
+        else:
+            spy = mocker.spy(qml.math, "einsum")
+
+        t = [
+            np.arange(4**8).reshape((4,) * 8),
+            np.arange(4**4).reshape((4,) * 4),
+            np.arange(4**2).reshape((4,) * 2),
+        ]
+        m = [
+            [
+                qcut.MeasureNode(wires=3),
+                qcut.MeasureNode(wires=1),
+                qcut.MeasureNode(wires=2),
+                qcut.MeasureNode(wires=3),
+                qcut.MeasureNode(wires=4),
+            ],
+            [
+                qcut.MeasureNode(wires=1),
+                qcut.MeasureNode(wires=2),
+            ],
+            [],
+        ]
+        p = [
+            [
+                qcut.PrepareNode(wires=1),
+                qcut.PrepareNode(wires=2),
+                qcut.PrepareNode(wires=3),
+            ],
+            [
+                qcut.PrepareNode(wires=1),
+                qcut.PrepareNode(wires=2),
+            ],
+            [
+                qcut.PrepareNode(wires=4),
+                qcut.PrepareNode(wires=3),
+            ],
+        ]
+        edges = [
+            (0, 0, 0, {"pair": (m[0][0], p[0][2])}),
+            (0, 1, 0, {"pair": (m[0][1], p[1][0])}),
+            (0, 1, 1, {"pair": (m[0][2], p[1][1])}),
+            (0, 2, 0, {"pair": (m[0][4], p[2][0])}),
+            (0, 2, 1, {"pair": (m[0][3], p[2][1])}),
+            (1, 0, 0, {"pair": (m[1][0], p[0][0])}),
+            (1, 0, 1, {"pair": (m[1][1], p[0][1])}),
+        ]
+        g = MultiDiGraph(edges)
+
+        res = qcut.contract_tensors(t, g, p, m, use_opt_einsum=use_opt_einsum)
+
+        eqn = spy.call_args[0][0]
+        expected_eqn = "abccdegf,deab,fg"
+
+        assert eqn == expected_eqn
+        assert np.allclose(res, np.einsum(eqn, *t))