Repository: PennyLaneAI/pennylane
Pull request #2257 (2022-03-01 17:05:58)

Circuit cutting: Remap fragment tape wires to match wires available on device
**Context:**

This PR ensures that N-qubit fragment circuits can be run on N-qubit devices by mapping the wires of the input circuit to match those available in the device.

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.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2257](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (f9b99a7) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/9c344f586f482e6b860c390152e91c74f4c2f6eb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (9c344f5) will **decrease** coverage by `2.66%`.
> The diff coverage is `95.45%`.

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

```diff
@@            Coverage Diff             @@
##           master    #2257      +/-   ##
==========================================
- Coverage   99.27%   96.61%   -2.67%     
==========================================
  Files         231      232       +1     
  Lines       18381    18734     +353     
==========================================
- Hits        18247    18099     -148     
- Misses        134      635     +501     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257/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% <ø> (ø)` | |
| [pennylane/transforms/qcut.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvcWN1dC5weQ==) | `99.76% <95.45%> (-0.24%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257/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/2257/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/2257/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/2257/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/2257/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/2257/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/2257/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/2257/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: |
| ... and [22 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257/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/2257?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/2257?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [9c344f5...f9b99a7](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2257?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-dev.md b/doc/releases/changelog-dev.md
index c67fc225c..fba9b319e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -211,6 +211,10 @@
   
   `WireCut` operations that do not lead to a disconnection are now being removed.
   [(#2260)](https://github.com/PennyLaneAI/pennylane/pull/2260)
+  
+  - Circuit cutting now remaps the wires of fragment circuits to match the available wires on the
+    device.
+    [(#2257)](https://github.com/PennyLaneAI/pennylane/pull/2257)
 
 <h3>Improvements</h3>
 
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 35ac77a12..29925f8b1 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -116,6 +116,7 @@ The following are utility functions that compose the circuit cutting transform.
     ~transforms.expand_fragment_tapes
     ~transforms.contract_tensors
     ~transforms.qcut_processing_fn
+    ~transforms.remap_tape_wires
     ~transforms.CutStrategy
 
 Transforms that act on tapes
@@ -208,5 +209,6 @@ from .qcut import (
     contract_tensors,
     qcut_processing_fn,
     cut_circuit,
+    remap_tape_wires,
     CutStrategy,
 )
diff --git a/pennylane/transforms/qcut.py b/pennylane/transforms/qcut.py
index 20b4a092c..d52e15f24 100644
--- a/pennylane/transforms/qcut.py
+++ b/pennylane/transforms/qcut.py
@@ -23,7 +23,7 @@ import warnings
 from dataclasses import InitVar, dataclass
 from functools import partial
 from itertools import product
-from typing import Any, Callable, ClassVar, Dict, List, Sequence, Tuple, Union
+from typing import Any, Callable, ClassVar, Dict, List, Optional, Sequence, Tuple, Union
 
 from networkx import MultiDiGraph, weakly_connected_components, has_path
 
@@ -902,7 +902,7 @@ def qcut_processing_fn(
 
 @batch_transform
 def cut_circuit(
-    tape: QuantumTape, use_opt_einsum: bool = False
+    tape: QuantumTape, use_opt_einsum: bool = False, device_wires: Optional[Wires] = None
 ) -> Tuple[Tuple[QuantumTape], Callable]:
     """
     Batch transform for circuit cutting.
@@ -919,10 +919,12 @@ def cut_circuit(
             for faster 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.
+        device_wires (.wires.Wires): Wires of the device that the cut circuits are to be run on
 
     Returns:
-        Tuple[Tuple[QuantumTape], Callable]: the tapes corresponding to the circuit fragments as a result of cutting
-        and a post-processing function which combines the results via tensor contractions.
+        Tuple[Tuple[QuantumTape], Callable]: the tapes corresponding to the circuit fragments as a
+        result of cutting and a post-processing function which combines the results via tensor
+        contractions.
 
     **Example**
 
@@ -992,6 +994,7 @@ def cut_circuit(
     replace_wire_cut_nodes(g)
     fragments, communication_graph = fragment_graph(g)
     fragment_tapes = [graph_to_tape(f) for f in fragments]
+    fragment_tapes = [remap_tape_wires(t, device_wires) for t in fragment_tapes]
     expanded = [expand_fragment_tapes(t) for t in fragment_tapes]
 
     configurations = []
@@ -1013,6 +1016,85 @@ def cut_circuit(
     )
 
 
+@cut_circuit.custom_qnode_wrapper
+def qnode_execution_wrapper(self, qnode, targs, tkwargs):
+    """Here, we overwrite the QNode execution wrapper in order
+    to access the device wires."""
+
+    tkwargs.setdefault("device_wires", qnode.device.wires)
+    return self.default_qnode_wrapper(qnode, targs, tkwargs)
+
+
+def remap_tape_wires(tape: QuantumTape, wires: Sequence) -> QuantumTape:
+    """Map the wires of a tape to a new set of wires.
+
+    Given an :math:`n`-wire ``tape``, this function returns a new :class:`~.QuantumTape` with
+    operations and measurements acting on the first :math:`n` wires provided in the ``wires``
+    argument. The input ``tape`` is left unmodified.
+
+    .. 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.
+
+    Args:
+        tape (QuantumTape): the quantum tape whose wires should be remapped
+        wires (Sequence): the new set of wires to map to
+
+    Returns:
+        QuantumTape: A remapped copy of the input tape
+
+    Raises:
+        ValueError: if the number of wires in ``tape`` exceeds ``len(wires)``
+
+    **Example**
+
+    .. code-block:: python
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.5, wires=2)
+            qml.RY(0.6, wires=3)
+            qml.CNOT(wires=[2, 3])
+            qml.expval(qml.PauliZ(2) @ qml.PauliZ(3))
+
+        new_wires = [0, 1]
+        new_tape = qml.transforms.remap_tape_wires(tape, new_wires)
+
+    >>> print(new_tape.draw())
+     0: ──RX(0.5)──╭C──╭┤ ⟨Z ⊗ Z⟩
+     1: ──RY(0.6)──╰X──╰┤ ⟨Z ⊗ Z⟩
+    """
+    if len(tape.wires) > len(wires):
+        raise ValueError(
+            f"Attempting to run a {len(tape.wires)}-wire circuit on a "
+            f"{len(wires)}-wire device. Consider increasing the number of wires in "
+            f"your device."
+        )
+
+    wire_map = dict(zip(tape.wires, wires))
+    copy_ops = [copy.copy(op) for op in tape.operations]
+    copy_meas = [copy.copy(op) for op in tape.measurements]
+
+    with QuantumTape() as new_tape:
+        for op in copy_ops:
+            new_wires = Wires([wire_map[w] for w in op.wires])
+            op._wires = new_wires
+            apply(op)
+        for meas in copy_meas:
+            obs = meas.obs
+
+            if isinstance(obs, Tensor):
+                for obs in obs.obs:
+                    new_wires = Wires([wire_map[w] for w in obs.wires])
+                    obs._wires = new_wires
+            else:
+                new_wires = Wires([wire_map[w] for w in obs.wires])
+                obs._wires = new_wires
+            apply(meas)
+
+    return new_tape
+
+
 @dataclass()
 class CutStrategy:
     """


Test patch:
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index 782c16aee..2c30bdb5b 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -2091,6 +2091,36 @@ class TestCutCircuitTransform:
 
         assert np.isclose(gradient, cut_gradient)
 
+    def test_device_wires(self, use_opt_einsum):
+        """Tests that a 3-qubit circuit is cut into two 2-qubit fragments such that both fragments
+        can be run on a 2-qubit device"""
+
+        def circuit():
+            qml.RX(0.4, wires=0)
+            qml.RX(0.5, wires=1)
+            qml.RX(0.6, wires=2)
+
+            qml.CNOT(wires=[0, 1])
+            qml.WireCut(wires=1)
+            qml.CNOT(wires=[1, 2])
+
+            return qml.expval(qml.PauliX(1) @ qml.PauliY(2))
+
+        dev_uncut = qml.device("default.qubit", wires=3)
+        dev_1 = qml.device("default.qubit", wires=2)
+        dev_2 = qml.device("default.qubit", wires=["Alice", 3.14, "Bob"])
+
+        uncut_circuit = qml.QNode(circuit, dev_uncut)
+        cut_circuit_1 = qml.transforms.cut_circuit(qml.QNode(circuit, dev_1), use_opt_einsum)
+        cut_circuit_2 = qml.transforms.cut_circuit(qml.QNode(circuit, dev_2), use_opt_einsum)
+
+        res_expected = uncut_circuit()
+        res_1 = cut_circuit_1()
+        res_2 = cut_circuit_2()
+
+        assert np.isclose(res_expected, res_1)
+        assert np.isclose(res_expected, res_2)
+
     def test_circuit_with_disconnected_components(self, use_opt_einsum, mocker):
         """Tests if a circuit that is fragmented into subcircuits such that some of the subcircuits
         are disconnected from the final terminal measurements is executed correctly"""
@@ -2115,6 +2145,59 @@ class TestCutCircuitTransform:
         assert spy.call_args[0][0][0].shape == ()
 
 
+class TestRemapTapeWires:
+    """Tests for the remap_tape_wires function"""
+
+    def test_raises(self):
+        """Test if a ValueError is raised when too few wires are provided"""
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.5, wires=2)
+            qml.RY(0.6, wires=3)
+            qml.CNOT(wires=[2, 3])
+            qml.expval(qml.PauliZ(2) @ qml.PauliZ(3))
+
+        with pytest.raises(ValueError, match="a 2-wire circuit on a 1-wire device"):
+            qcut.remap_tape_wires(tape, [0])
+
+    def test_mapping(self):
+        """Test if the function returns the expected tape when an observable measurement is
+        used"""
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.5, wires=2)
+            qml.RY(0.6, wires=3)
+            qml.CNOT(wires=[2, 3])
+            qml.expval(qml.PauliZ(2))
+
+        with qml.tape.QuantumTape() as expected_tape:
+            qml.RX(0.5, wires=0)
+            qml.RY(0.6, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0))
+
+        new_tape = qcut.remap_tape_wires(tape, [0, 1])
+
+        compare_tapes(expected_tape, new_tape)
+
+    def test_mapping_tensor(self):
+        """Test if the function returns the expected tape when a tensor product measurement is
+        used"""
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(0.5, wires=2)
+            qml.RY(0.6, wires=3)
+            qml.CNOT(wires=[2, 3])
+            qml.expval(qml.PauliZ(2) @ qml.PauliZ(3))
+
+        with qml.tape.QuantumTape() as expected_tape:
+            qml.RX(0.5, wires=0)
+            qml.RY(0.6, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        new_tape = qcut.remap_tape_wires(tape, [0, 1])
+
+        compare_tapes(expected_tape, new_tape)
+
+
 class TestCutCircuitTransformValidation:
     """Tests of validation checks in the cut_circuit function"""