Repository: PennyLaneAI/pennylane
Pull request #7340 (2025-04-28 18:10:04)

BugFix: fix special case discontinuity in two-qubit decomposition
Follow up to https://github.com/PennyLaneAI/pennylane/pull/5448, applying the same solution to the 3-cnot case.

Fixes: https://github.com/PennyLaneAI/pennylane/issues/7339
[sc-90052]

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://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7340?dropdown=coverage&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:
> Project coverage is 99.65%. Comparing base [(`067543f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/067543f08d38950b5db5e7817c4a9c9931a75435?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`9d9434e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/9d9434e41b364a3450018e4d56f9f2948555f5d9?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

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


```diff
@@           Coverage Diff           @@
##           master    #7340   +/-   ##
=======================================
  Coverage   99.65%   99.65%           
=======================================
  Files         526      526           
  Lines       50291    50292    +1     
=======================================
+ Hits        50115    50116    +1     
  Misses        176      176           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7340?dropdown=coverage&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).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 83547767c..a505e4bb8 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -234,6 +234,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* A fix was made to `default.qubit` to allow for using `qml.Snapshot` with defer-measurements (`mcm_method="deferred"`).
+  [(#7335)](https://github.com/PennyLaneAI/pennylane/pull/7335)
+
 * Fixes the repr for empty `Prod` and `Sum` instances to better communicate the existence of an empty instance.
   [(#7346)](https://github.com/PennyLaneAI/pennylane/pull/7346)
 
@@ -267,6 +270,9 @@
 * Fixed a bug where the phase is used as the wire label for a `qml.GlobalPhase` when capture is enabled.
   [(#7211)](https://github.com/PennyLaneAI/pennylane/pull/7211)
 
+* Fixed a bug where `two_qubit_decomposition` provides an incorrect decomposition for some special matrices.
+  [(#7340)](https://github.com/PennyLaneAI/pennylane/pull/7340)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 859860663..a941db7a3 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -560,10 +560,11 @@ class DefaultQubit(Device):
 
         if config.interface == qml.math.Interface.JAX_JIT:
             transform_program.add_transform(no_counts)
-        transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
         transform_program.add_transform(
             mid_circuit_measurements, device=self, mcm_config=config.mcm_config
         )
+        # validate_device_wires needs to be after defer_measurement has added more wires.
+        transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
         transform_program.add_transform(
             decompose,
             stopping_condition=stopping_condition,
diff --git a/pennylane/ops/op_math/decompositions/unitary_decompositions.py b/pennylane/ops/op_math/decompositions/unitary_decompositions.py
index 7167e17e2..ee23839b3 100644
--- a/pennylane/ops/op_math/decompositions/unitary_decompositions.py
+++ b/pennylane/ops/op_math/decompositions/unitary_decompositions.py
@@ -679,8 +679,9 @@ def _decompose_3_cnots(U, wires, initial_phase):
     # -╭V- = -╭X--RZ(d)--╭C---------╭X--╭SWAP-
     # -╰V- = -╰C--RY(b)--╰X--RY(a)--╰C--╰SWAP-
 
-    RZd = ops.RZ.compute_matrix(math.cast_like(delta, 1j))
-    RYb = ops.RY.compute_matrix(beta)
+    EPS = math.finfo(delta.dtype).eps
+    RZd = ops.RZ.compute_matrix(math.cast_like(delta + 5 * EPS, 1j))
+    RYb = ops.RY.compute_matrix(beta + 5 * EPS)
     RYa = ops.RY.compute_matrix(alpha)
 
     V_mats = [


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index 879636d31..79341835a 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -2215,6 +2215,22 @@ class TestIntegration:
 
         assert qml.math.allclose(grad, grad_jit)
 
+    def test_snapshot_with_defer_measurement(self):
+        """Test that snapshots can be taken with defer_measurements."""
+
+        dev = qml.device("default.qubit")
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=0)
+            qml.measure(0)
+            qml.Snapshot("label")
+            return qml.probs(wires=0)
+
+        snapshots = qml.snapshots(func)()
+        assert snapshots["label"].shape == (4,)
+        assert qml.math.allclose(snapshots["execution_results"], np.array([0.5, 0.5]))
+
 
 @pytest.mark.parametrize("max_workers", max_workers_list)
 def test_broadcasted_parameter(max_workers):
diff --git a/tests/ops/functions/test_iterative_qpe.py b/tests/ops/functions/test_iterative_qpe.py
index 515044d0d..110d5a53e 100644
--- a/tests/ops/functions/test_iterative_qpe.py
+++ b/tests/ops/functions/test_iterative_qpe.py
@@ -267,7 +267,7 @@ def test_capture_execution(seed):
 
     jaxpr = jax.make_jaxpr(f)(1.5)
 
-    dev = qml.device("default.qubit", wires=3, seed=seed)
+    dev = qml.device("default.qubit", wires=5, seed=seed)
 
     # hack for single-branch statistics
     samples = qml.math.vstack([dev.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, x) for _ in range(5000)])
diff --git a/tests/ops/op_math/test_decompositions.py b/tests/ops/op_math/test_decompositions.py
index 750122c7d..6447f5cc5 100644
--- a/tests/ops/op_math/test_decompositions.py
+++ b/tests/ops/op_math/test_decompositions.py
@@ -1210,38 +1210,52 @@ class TestTwoQubitUnitaryDecompositionInterfaces:
         assert check_matrix_equivalence(U, jitted_matrix, atol=1e-7)
 
 
-def test_two_qubit_decomposition_special_case_discontinuity():
-    """Test that two_qubit_decomposition still provides accurate numbers at a special case."""
-
-    def make_unitary(theta1):
-        generator = (
-            theta1
+# This was the routine used to generate the problematic matrix in the original
+# bug report (see https://github.com/PennyLaneAI/pennylane/issues/5308)
+def _make_unitary(theta1):
+    generator = (
+        theta1
+        / 2
+        * (
+            np.cos(0.2) / 2 * (np.array([[0, 0, 0, 0], [0, 0, 2, 0], [0, 2, 0, 0], [0, 0, 0, 0]]))
+            + np.sin(0.2)
             / 2
-            * (
-                np.cos(0.2)
-                / 2
-                * (np.array([[0, 0, 0, 0], [0, 0, 2, 0], [0, 2, 0, 0], [0, 0, 0, 0]]))
-                + np.sin(0.2)
-                / 2
-                * (np.array([[0, 0, 0, 0], [0, 0, 2j, 0], [0, -2j, 0, 0], [0, 0, 0, 0]]))
-            )
+            * (np.array([[0, 0, 0, 0], [0, 0, 2j, 0], [0, -2j, 0, 0], [0, 0, 0, 0]]))
         )
+    )
 
-        def expm(val):
-            d, U = np.linalg.eigh(-1.0j * val)
-            return np.dot(U, np.dot(np.diag(np.exp(1.0j * d)), np.conj(U).T))
+    def expm(val):
+        d, U = np.linalg.eigh(-1.0j * val)
+        return np.dot(U, np.dot(np.diag(np.exp(1.0j * d)), np.conj(U).T))
 
-        mat = expm(-1j * generator)
+    mat = expm(-1j * generator)
 
-        assert np.allclose(
-            np.dot(np.transpose(np.conj(mat)), mat), np.eye(len(mat))
-        ), "mat is not unitary"
+    assert np.allclose(
+        np.dot(np.transpose(np.conj(mat)), mat), np.eye(len(mat))
+    ), "mat is not unitary"
 
-        return mat
+    return mat
+
+
+@pytest.mark.parametrize(
+    "U",
+    [
+        _make_unitary(np.pi / 2),
+        np.array(
+            [
+                [-1, 0, 0, 0],
+                [0, 1 / np.sqrt(2), 1 / np.sqrt(2), 0],
+                [0, 1 / np.sqrt(2), -1 / np.sqrt(2), 0],
+                [0, 0, 0, -1],
+            ]
+        ),
+    ],
+)
+def test_two_qubit_decomposition_special_case_discontinuity(U):
+    """Test that two_qubit_decomposition still provides accurate numbers at a special case."""
 
-    mat = make_unitary(np.pi / 2)
-    decomp_mat = qml.matrix(two_qubit_decomposition, wire_order=(0, 1))(mat, wires=(0, 1))
-    assert qml.math.allclose(mat, decomp_mat)
+    decomp_mat = qml.matrix(two_qubit_decomposition, wire_order=(0, 1))(U, wires=(0, 1))
+    assert qml.math.allclose(U, decomp_mat)
 
 
 class TestTwoQubitDecompositionWarnings:
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index 9fb0961d4..59831a7ce 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -577,7 +577,7 @@ class TestPyTreeStructure:
     )
     def test_tensor_measurement(self, measurement):
         """Tests that measurements of tensor type are handled correctly"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit", wires=3)
 
         @qml.qnode(dev)
         def circuit():
diff --git a/tests/transforms/test_add_noise.py b/tests/transforms/test_add_noise.py
index fd069a8f6..902eef7f3 100644
--- a/tests/transforms/test_add_noise.py
+++ b/tests/transforms/test_add_noise.py
@@ -282,7 +282,7 @@ class TestAddNoiseInterface:
     def test_add_noise_with_non_qwc_obs_and_mid_meas(self):
         """Test that the add_noise transform catches and reports errors from the enclosed function."""
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit", wires=5)
 
         fcond = qml.noise.wires_in([0, 1])
 
diff --git a/tests/workflow/test_construct_batch.py b/tests/workflow/test_construct_batch.py
index d855a6323..1cf535adc 100644
--- a/tests/workflow/test_construct_batch.py
+++ b/tests/workflow/test_construct_batch.py
@@ -114,7 +114,7 @@ class TestTransformProgramGetter:
         p_sliced = get_transform_program(circuit, slice(2, 7, 2))
         assert len(p_sliced) == 3
         assert p_sliced[0].transform == qml.compile.transform
-        assert p_sliced[1].transform == qml.devices.preprocess.validate_device_wires.transform
+        assert p_sliced[1].transform == qml.devices.preprocess.mid_circuit_measurements.transform
         assert p_sliced[2].transform == qml.devices.preprocess.decompose.transform
 
     def test_diff_method_device_gradient(self):