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]>

.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)
 

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))

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)