Repository: PennyLaneAI/pennylane
Pull request #4732 (2023-10-25 18:53:08)

fix transpile for state measurement
Fixes Issue #4723 [sc-48555]


Patch:
diff --git a/doc/releases/changelog-0.33.0.md b/doc/releases/changelog-0.33.0.md
index 553087421..33f3de25d 100644
--- a/doc/releases/changelog-0.33.0.md
+++ b/doc/releases/changelog-0.33.0.md
@@ -538,6 +538,10 @@
 <h3>Bug fixes 🐛</h3>
 
 * Fixes `LocalHilbertSchmidt.compute_decomposition` so the template can be used in a qnode.
+  [(#4719)](https://github.com/PennyLaneAI/pennylane/pull/4719)
+
+* Fixes `transforms.transpile` with arbitrary measurement processes.
+  [(#4732)](https://github.com/PennyLaneAI/pennylane/pull/4732)
 
 * Providing `work_wires=None` to `qml.GroverOperator` no longer interprets `None` as a wire.
   [(#4668)](https://github.com/PennyLaneAI/pennylane/pull/4668)
diff --git a/pennylane/transforms/transpile.py b/pennylane/transforms/transpile.py
index 5999e9a8c..cd09e47d3 100644
--- a/pennylane/transforms/transpile.py
+++ b/pennylane/transforms/transpile.py
@@ -12,7 +12,6 @@ from pennylane.ops import __all__ as all_ops
 from pennylane.ops.qubit import SWAP
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumTape
-from pennylane.wires import Wires
 
 
 @transform
@@ -107,8 +106,6 @@ def transpile(
             "The transpile transform only supports gates acting on 1 or 2 qubits."
         )
 
-    gates = []
-
     # we wrap all manipulations inside stop_recording() so that we don't queue anything due to unrolling of templates
     # or newly applied swap gates
     with QueuingManager.stop_recording():
@@ -121,6 +118,7 @@ def transpile(
         # make copy of ops
         list_op_copy = expanded_tape.operations.copy()
         measurements = expanded_tape.measurements.copy()
+        gates = []
 
         while len(list_op_copy) > 0:
             op = list_op_copy[0]
@@ -141,7 +139,7 @@ def transpile(
                 continue
 
             # since in each iteration, we adjust indices of each op, we reset logical -> phyiscal mapping
-            map_wires = {w: w for w in tape.wires}
+            wire_map = {w: w for w in tape.wires}
 
             # to make sure two qubit gates which act on non-neighbouring qubits q1, q2 can be applied, we first look
             # for the shortest path between the two qubits in the connectivity graph. We then move the q2 into the
@@ -156,17 +154,16 @@ def transpile(
                 # swap wires
                 gates.append(SWAP(wires=[w0, w1]))
                 # update logical -> phyiscal mapping
-                map_wires = {
-                    k: (w0 if v == w1 else (w1 if v == w0 else v)) for k, v in map_wires.items()
+                wire_map = {
+                    k: (w0 if v == w1 else (w1 if v == w0 else v)) for k, v in wire_map.items()
                 }
 
             # append op to gates with adjusted indices and remove from list
-            gates.append(_adjust_op_indices(op, map_wires))
+            gates.append(op.map_wires(wire_map))
             list_op_copy.pop(0)
 
-            # adjust qubit indices in remaining ops + measurements to new mapping
-            list_op_copy = [_adjust_op_indices(op, map_wires) for op in list_op_copy]
-            measurements = [_adjust_mmt_indices(m, map_wires) for m in measurements]
+            list_op_copy = [op.map_wires(wire_map) for op in list_op_copy]
+            measurements = [m.map_wires(wire_map) for m in measurements]
     new_tape = type(tape)(gates, measurements, shots=tape.shots)
 
     def null_postprocessing(results):
@@ -176,24 +173,3 @@ def transpile(
         return results[0]
 
     return [new_tape], null_postprocessing
-
-
-def _adjust_op_indices(_op, _map_wires):
-    """helper function which adjusts wires in Operation according to the map _map_wires"""
-    _new_wires = Wires([_map_wires[w] for w in _op.wires])
-    _params = _op.parameters
-    if len(_params) == 0:
-        return type(_op)(wires=_new_wires)
-    return type(_op)(*_params, wires=_new_wires)
-
-
-def _adjust_mmt_indices(_m, _map_wires):
-    """helper function which adjusts wires in MeasurementProcess according to the map _map_wires"""
-    _new_wires = Wires([_map_wires[w] for w in _m.wires])
-
-    # change wires of observable
-    if _m.obs is None:
-        return type(_m)(eigvals=_m.eigvals(), wires=_new_wires)
-
-    _new_obs = type(_m.obs)(wires=_new_wires, id=_m.obs.id)
-    return type(_m)(obs=_new_obs)


Test patch:
diff --git a/tests/transforms/test_transpile.py b/tests/transforms/test_transpile.py
index 438aa5a05..e9708a9ea 100644
--- a/tests/transforms/test_transpile.py
+++ b/tests/transforms/test_transpile.py
@@ -295,3 +295,18 @@ class TestTranspile:
         assert qml.math.allclose(
             original_expectation, transpiled_expectation, atol=np.finfo(float).eps
         )
+
+    def test_transpile_state(self):
+        """Test that transpile works with state measurement process."""
+
+        tape = qml.tape.QuantumScript([qml.PauliX(0), qml.CNOT((0, 2))], [qml.state()], shots=100)
+        batch, fn = qml.transforms.transpile(tape, coupling_map=[(0, 1), (1, 2)])
+
+        assert len(batch) == 1
+        assert fn(("a",)) == "a"
+
+        assert batch[0][0] == qml.PauliX(0)
+        assert batch[0][1] == qml.SWAP((1, 2))
+        assert batch[0][2] == qml.CNOT((0, 1))
+        assert batch[0][3] == qml.state()
+        assert batch[0].shots == tape.shots