Update serialization tests with PauliWord order updates (#688)

* port PR changes so far, after traumatic rebase * update vjp tests * format * update vjp tests * update dev version * format * add tests for vjp * Auto update version * Auto update version * Updated tests * Updated serialization * Unpinned cmake * Update changelog * Update .github/CHANGELOG.md * isort * black * Run CI with updated PL branch * Auto update version * trigger ci * Added changes per review * Auto update version * Update requirements-dev.txt * Update pennylane_lightning/core/_serialize.py Co-authored-by: Ali Asadi <[email protected]> * Auto update version * Trigger CI * Update overlapping wires prod serialization test * Updated tests to work with new pauli word order * Added device_vjp tests * Update serialize tests * Auto update version * Skipped tf test * Auto update version * Auto update version * Trigger CI * Updated torch test * Update serialization tests; removed xfail * Trigger CI * Auto update version * Remove vjp tests * Trigger CI * Trigger CI * [skip ci] Revert requirements file * Auto update version from '0.36.0-dev35' to '0.36.0-dev36' * Trigger CI * Trigger CI * Fix native mcm workflow following dynamic_one_shot refactor. (#694) * Fix native mcm workflow following dynamic_one_shot refactor. * Auto update version from '0.36.0-dev35' to '0.36.0-dev36' * Update changelog. --------- Co-authored-by: ringo-but-quantum <[email protected]> --------- Co-authored-by: AmintorDusko <[email protected]> Co-authored-by: Amintor Dusko <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]> Co-authored-by: Ali Asadi <[email protected]> Co-authored-by: ringo-but-quantum <[email protected]> Co-authored-by: Vincent Michaud-Rioux <[email protected]>
PennyLaneAI · Apr 24, 2024 · 6efe874 · 6efe874
1 parent 43bed84
commit 6efe874
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 30 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -86,6 +86,9 @@
 
 ### Bug fixes
 
+* `dynamic_one_shot` was refactored to use `SampleMP` measurements as a way to return the mid-circuit measurement samples. `LightningQubit`'s `simulate` is modified accordingly.
+  [(#694)](https://github.com/PennyLaneAI/pennylane/pull/694)
+
 * `LightningQubit` correctly decomposes state prep operations when used in the middle of a circuit.
   [(#687)](https://github.com/PennyLaneAI/pennylane/pull/687)
 

diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev35"
+__version__ = "0.36.0-dev36"
diff --git a/pennylane_lightning/lightning_qubit/_measurements.py b/pennylane_lightning/lightning_qubit/_measurements.py
@@ -248,14 +248,15 @@ def measurement(self, measurementprocess: MeasurementProcess) -> TensorLike:
         """
         return self.get_measurement_function(measurementprocess)(measurementprocess)
 
-    def measure_final_state(self, circuit: QuantumScript) -> Result:
+    def measure_final_state(self, circuit: QuantumScript, mid_measurements=None) -> Result:
         """
         Perform the measurements required by the circuit on the provided state.
 
         This is an internal function that will be called by the successor to ``lightning.qubit``.
 
         Args:
             circuit (QuantumScript): The single circuit to simulate
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
 
         Returns:
             Tuple[TensorLike]: The measurement results
@@ -272,6 +273,7 @@ def measure_final_state(self, circuit: QuantumScript) -> Result:
         results = self.measure_with_samples(
             circuit.measurements,
             shots=circuit.shots,
+            mid_measurements=mid_measurements,
         )
 
         if len(circuit.measurements) == 1:
@@ -285,27 +287,37 @@ def measure_final_state(self, circuit: QuantumScript) -> Result:
     # pylint:disable = too-many-arguments
     def measure_with_samples(
         self,
-        mps: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]],
+        measurements: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]],
         shots: Shots,
+        mid_measurements=None,
     ) -> List[TensorLike]:
         """
         Returns the samples of the measurement process performed on the given state.
         This function assumes that the user-defined wire labels in the measurement process
         have already been mapped to integer wires used in the device.
 
         Args:
-            mps (List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]]):
+            measurements (List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]]):
                 The sample measurements to perform
             shots (Shots): The number of samples to take
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
 
         Returns:
             List[TensorLike[Any]]: Sample measurement results
         """
+        # last N measurements are sampling MCMs in ``dynamic_one_shot`` execution mode
+        mps = measurements[0 : -len(mid_measurements)] if mid_measurements else measurements
+        skip_measure = (
+            any(v == -1 for v in mid_measurements.values()) if mid_measurements else False
+        )
 
         groups, indices = _group_measurements(mps)
 
         all_res = []
         for group in groups:
+            if skip_measure:
+                all_res.extend([None] * len(group))
+                continue
             if isinstance(group[0], (ExpectationMP, VarianceMP)) and isinstance(
                 group[0].obs, SparseHamiltonian
             ):
@@ -333,6 +345,10 @@ def measure_with_samples(
             res for _, res in sorted(list(enumerate(all_res)), key=lambda r: flat_indices[r[0]])
         )
 
+        # append MCM samples
+        if mid_measurements:
+            sorted_res += tuple(mid_measurements.values())
+
         # put the shot vector axis before the measurement axis
         if shots.has_partitioned_shots:
             sorted_res = tuple(zip(*sorted_res))

diff --git a/pennylane_lightning/lightning_qubit/lightning_qubit.py b/pennylane_lightning/lightning_qubit/lightning_qubit.py
@@ -80,11 +80,8 @@ def simulate(circuit: QuantumScript, state: LightningStateVector, mcmc: dict = N
     if circuit.shots and has_mcm:
         mid_measurements = {}
         final_state = state.get_final_state(circuit, mid_measurements=mid_measurements)
-        if any(v == -1 for v in mid_measurements.values()):
-            return None, mid_measurements
-        return (
-            LightningMeasurements(final_state, **mcmc).measure_final_state(circuit),
-            mid_measurements,
+        return LightningMeasurements(final_state, **mcmc).measure_final_state(
+            circuit, mid_measurements=mid_measurements
         )
     final_state = state.get_final_state(circuit)
     return LightningMeasurements(final_state, **mcmc).measure_final_state(circuit)

diff --git a/tests/test_serialize.py b/tests/test_serialize.py
@@ -146,10 +146,9 @@ def test_tensor_non_tensor_return(self, use_csingle, wires_map):
 
         named_obs = NamedObsC64 if use_csingle else NamedObsC128
         tensor_prod_obs = TensorProdObsC64 if use_csingle else TensorProdObsC128
-        first_s = tensor_prod_obs([named_obs("PauliX", [1]), named_obs("PauliZ", [0])])
 
         s_expected = [
-            first_s,
+            tensor_prod_obs([named_obs("PauliZ", [0]), named_obs("PauliX", [1])]),
             named_obs("Hadamard", [1]),
         ]
 
@@ -158,7 +157,6 @@ def test_tensor_non_tensor_return(self, use_csingle, wires_map):
         )
         assert s == s_expected
 
-    @pytest.mark.xfail(reason="Prod with overlapping wires not supported")
     @pytest.mark.parametrize("use_csingle", [True, False])
     @pytest.mark.parametrize("wires_map", [wires_dict, None])
     def test_prod_return_with_overlapping_wires(self, use_csingle, wires_map):
@@ -283,7 +281,7 @@ def test_hamiltonian_return(self, use_csingle, wires_map):
                         named_obs("PauliY", [2]),
                     ]
                 ),
-                tensor_prod_obs([named_obs("PauliY", [2]), named_obs("PauliX", [0])]),
+                tensor_prod_obs([named_obs("PauliX", [0]), named_obs("PauliY", [2])]),
                 hermitian_obs(np.ones(64, dtype=c_dtype), [0, 1, 2]),
             ],
         )
@@ -317,22 +315,23 @@ def test_hamiltonian_tensor_return(self, use_csingle, wires_map):
             tape, wires_map
         )
 
-        first_term = tensor_prod_obs(
-            [
-                hermitian_obs(np.eye(4, dtype=c_dtype).ravel(), [0, 1]),
-                named_obs("PauliY", [2]),
-                named_obs("PauliZ", [3]),
-            ]
-        )
+        # Expression (ham @ obs) is converted internally by Pennylane
+        # where obs is appended to each term of the ham
 
         s_expected = hamiltonian_obs(
             np.array([0.3, 0.5, 0.4], dtype=r_dtype),
             [
-                first_term,
                 tensor_prod_obs(
                     [
+                        hermitian_obs(np.eye(4, dtype=c_dtype).ravel(), [0, 1]),
                         named_obs("PauliY", [2]),
+                        named_obs("PauliZ", [3]),
+                    ]
+                ),
+                tensor_prod_obs(
+                    [
                         named_obs("PauliX", [0]),
+                        named_obs("PauliY", [2]),
                         named_obs("PauliZ", [3]),
                     ]
                 ),
@@ -388,7 +387,7 @@ def test_hamiltonian_mix_return(self, use_csingle, wires_map):
                         named_obs("PauliY", [2]),
                     ]
                 ),
-                tensor_prod_obs([named_obs("PauliY", [2]), named_obs("PauliX", [0])]),
+                tensor_prod_obs([named_obs("PauliX", [0]), named_obs("PauliY", [2])]),
                 hermitian_obs(np.ones(64, dtype=c_dtype), [0, 1, 2]),
             ],
         )
@@ -401,7 +400,7 @@ def test_hamiltonian_mix_return(self, use_csingle, wires_map):
                         hermitian_obs(np.eye(4, dtype=c_dtype).ravel(), [1, 2]),
                     ]
                 ),
-                tensor_prod_obs([named_obs("PauliX", [2]), named_obs("PauliY", [0])]),
+                tensor_prod_obs([named_obs("PauliY", [0]), named_obs("PauliX", [2])]),
             ],
         )
 
@@ -441,7 +440,7 @@ def test_pauli_rep_single_term(self, use_csingle, wires_map):
         s, _ = QuantumScriptSerializer(device_name, use_csingle).serialize_observables(
             tape, wires_map
         )
-        s_expected = tensor_prod_obs([named_obs("PauliZ", [1]), named_obs("PauliX", [0])])
+        s_expected = tensor_prod_obs([named_obs("PauliX", [0]), named_obs("PauliZ", [1])])
         assert s[0] == s_expected
 
     @pytest.mark.parametrize("use_csingle", [True, False])
@@ -482,11 +481,7 @@ def test_prod(self, use_csingle, wires_map):
         assert isinstance(res[0], tensor_obs)
 
         s_expected = tensor_obs(
-            [
-                named_obs("PauliZ", [0]),
-                named_obs("PauliX", [1]),
-                named_obs("Hadamard", [2]),
-            ]
+            [named_obs("PauliZ", [0]), named_obs("PauliX", [1]), named_obs("Hadamard", [2])]
         )
         assert res[0] == s_expected
 
@@ -522,7 +517,7 @@ def test_sum(self, use_csingle, wires_map):
             coeffs,
             [
                 tensor_obs(
-                    [named_obs("PauliZ", [1]), named_obs("PauliX", [0]), named_obs("PauliX", [2])]
+                    [named_obs("PauliX", [0]), named_obs("PauliZ", [1]), named_obs("PauliX", [2])]
                 ),
                 tensor_obs(
                     [named_obs("PauliZ", [0]), named_obs("PauliY", [1]), named_obs("Hadamard", [2])]