Repository: PennyLaneAI/pennylane
Pull request #5081 (2024-01-18 19:37:37)

fix `undo_swaps` and `split_non_commuting`
When exploring in PR #5058, I found bugs in `undo_swaps` and `split_non_commuting`.

[sc-54373]
[sc-54374]

Undo swaps had a fragile way of mapping wires, as it was created before we had `Operator.map_wires`.  Now it works with more general operations.

`split_non_commuting` was not passing the shots along. This was a silent issue on the old device interface, but would be causing problems if it was ever used with the new device interface.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5081?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`c01cb4b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/c01cb4b96f758a9f1a0c4e422115829b80e113cb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67% compared to head [(`74ed30f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5081?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67%.
> Report is 27 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5081      +/-   ##
==========================================
- Coverage   99.67%   99.67%   -0.01%     
==========================================
  Files         394      394              
  Lines       35670    35448     -222     
==========================================
- Hits        35554    35332     -222     
  Misses        116      116              
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5081?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?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 393882590..f035467ce 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -135,6 +135,12 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.transforms.undo_swaps` can now work with operators with hyperparameters or nesting.
+  [(#5081)](https://github.com/PennyLaneAI/pennylane/pull/5081)
+
+* `qml.transforms.split_non_commuting` will now pass the original shots along.
+  [(#5081)](https://github.com/PennyLaneAI/pennylane/pull/5081)
+
 * If `argnum` is provided to a gradient transform, only the parameters specified in `argnum` will have their gradient methods validated.
   [(#5035)](https://github.com/PennyLaneAI/pennylane/pull/5035)
 
diff --git a/pennylane/transforms/optimization/undo_swaps.py b/pennylane/transforms/optimization/undo_swaps.py
index 2f8dbcf3d..20c180330 100644
--- a/pennylane/transforms/optimization/undo_swaps.py
+++ b/pennylane/transforms/optimization/undo_swaps.py
@@ -19,8 +19,13 @@ from typing import Sequence, Callable
 from pennylane.transforms import transform
 
 from pennylane.tape import QuantumTape
-from pennylane.wires import Wires
-from pennylane.queuing import QueuingManager
+
+
+def null_postprocessing(results):
+    """A postprocesing function returned by a transform that only converts the batch of results
+    into a result for a single ``QuantumTape``.
+    """
+    return results[0]
 
 
 @transform
@@ -89,48 +94,24 @@ def undo_swaps(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):
             2: ──X──┤
 
     """
-    # Make a working copy of the list to traverse
-    list_copy = tape.operations.copy()
-    list_copy.reverse()
 
-    map_wires = {wire: wire for wire in tape.wires}
+    wire_map = {wire: wire for wire in tape.wires}
     gates = []
 
-    def _change_wires(wires):
-        change_wires = Wires([])
-        wires = wires.toarray()
-        for wire in wires:
-            change_wires += map_wires[wire]
-        return change_wires
-
-    with QueuingManager.stop_recording():
-        while len(list_copy) > 0:
-            current_gate = list_copy[0]
-            params = current_gate.parameters
-            if current_gate.name != "SWAP":
-                if len(params) == 0:
-                    gates.append(type(current_gate)(wires=_change_wires(current_gate.wires)))
-                else:
-                    gates.append(
-                        type(current_gate)(*params, wires=_change_wires(current_gate.wires))
-                    )
-
-            else:
-                swap_wires_0, swap_wires_1 = current_gate.wires
-                map_wires[swap_wires_0], map_wires[swap_wires_1] = (
-                    map_wires[swap_wires_1],
-                    map_wires[swap_wires_0],
-                )
-            list_copy.pop(0)
-
-        gates.reverse()
-
-    new_tape = type(tape)(gates, tape.measurements, shots=tape.shots)
-
-    def null_postprocessing(results):
-        """A postprocesing function returned by a transform that only converts the batch of results
-        into a result for a single ``QuantumTape``.
-        """
-        return results[0]
+    for current_gate in reversed(tape.operations):
+        if current_gate.name == "SWAP":
+            swap_wires_0, swap_wires_1 = current_gate.wires
+            wire_map[swap_wires_0], wire_map[swap_wires_1] = (
+                wire_map[swap_wires_1],
+                wire_map[swap_wires_0],
+            )
+        else:
+            gates.append(current_gate.map_wires(wire_map))
+
+    gates.reverse()
+
+    new_tape = type(tape)(
+        gates, tape.measurements, shots=tape.shots, trainable_params=tape.trainable_params
+    )
 
     return [new_tape], null_postprocessing
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index 3dec7b19c..cbc78deac 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -23,6 +23,13 @@ import pennylane as qml
 from pennylane.transforms import transform
 
 
+def null_postprocessing(results):
+    """A postprocesing function returned by a transform that only converts the batch of results
+    into a result for a single ``QuantumTape``.
+    """
+    return results[0]
+
+
 @transform
 def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
     r"""
@@ -194,8 +201,7 @@ def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.Quantu
         tapes = []
         for indices in group_coeffs:
             new_tape = tape.__class__(
-                tape.operations,
-                (tape.measurements[i] for i in indices),
+                tape.operations, (tape.measurements[i] for i in indices), shots=tape.shots
             )
 
             tapes.append(new_tape)
@@ -233,4 +239,4 @@ def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.Quantu
         return tapes, reorder_fn
 
     # if the group is already commuting, no need to do anything
-    return [tape], lambda res: res[0]
+    return [tape], null_postprocessing


Test patch:
diff --git a/tests/transforms/test_optimization/test_undo_swaps.py b/tests/transforms/test_optimization/test_undo_swaps.py
index 595bb51b9..a5ff7ec33 100644
--- a/tests/transforms/test_optimization/test_undo_swaps.py
+++ b/tests/transforms/test_optimization/test_undo_swaps.py
@@ -26,6 +26,32 @@ from pennylane.transforms.optimization import undo_swaps
 class TestUndoSwaps:
     """Test that check the main functionalities of the `undo_swaps` transform"""
 
+    def test_transform_non_standard_operations(self):
+        """Test that the transform works on non-standard operations with nesting or hyperparameters."""
+
+        ops = [
+            qml.adjoint(qml.S(0)),
+            qml.PauliRot(1.2, "XY", wires=(0, 2)),
+            qml.ctrl(qml.PauliX(0), 2, control_values=[0, 0]),
+            qml.SWAP((0, 1)),
+        ]
+
+        tape = qml.tape.QuantumScript(ops, [qml.state()], shots=100)
+        batch, fn = qml.transforms.undo_swaps(tape)
+
+        expected_ops = [
+            qml.adjoint(qml.S(1)),
+            qml.PauliRot(1.2, "XY", wires=(1, 2)),
+            qml.ctrl(qml.PauliX(1), 2, control_values=[0, 0]),
+        ]
+        assert len(batch) == 1
+        assert batch[0].shots == tape.shots
+
+        assert fn(["a"]) == "a"
+
+        for op1, expected in zip(batch[0].operations, expected_ops):
+            assert op1 == expected
+
     def test_one_qubit_gates_transform(self):
         """Test that a single-qubit gate changes correctly with a SWAP."""
 
diff --git a/tests/transforms/test_split_non_commuting.py b/tests/transforms/test_split_non_commuting.py
index 4b2ae1d73..9c68c0e27 100644
--- a/tests/transforms/test_split_non_commuting.py
+++ b/tests/transforms/test_split_non_commuting.py
@@ -64,8 +64,10 @@ class TestUnittestSplitNonCommuting:
             meas_type(op=qml.PauliZ(0) @ qml.PauliZ(3))
 
         # test transform on tape
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=100)
         split, fn = split_non_commuting(tape)
+        for t in split:
+            assert t.shots == tape.shots
 
         spy = mocker.spy(qml.math, "concatenate")
 
@@ -74,8 +76,10 @@ class TestUnittestSplitNonCommuting:
         assert fn([0.5]) == 0.5
 
         # test transform on qscript
-        qs = qml.tape.QuantumScript(tape.operations, tape.measurements)
+        qs = qml.tape.QuantumScript(tape.operations, tape.measurements, shots=50)
         split, fn = split_non_commuting(qs)
+        for t in split:
+            assert t.shots == qs.shots
 
         assert len(split) == 1
         assert all(isinstance(i_qs, qml.tape.QuantumScript) for i_qs in split)