Add MCM support in LQ (new device API)

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* `lightning.qubit` supports mid-circuit measurements.
+  [(#650)](https://github.com/PennyLaneAI/pennylane-lightning/pull/650)
+
 * Add finite shots support in `lightning.qubit2`.
   [(#630)](https://github.com/PennyLaneAI/pennylane-lightning/pull/630)
 

.github/workflows/tests_linux.yml

@@ -33,7 +33,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: C++ tests
     runs-on: ${{ matrix.os }}
 
@@ -106,7 +106,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: Python tests
     runs-on: ${{ matrix.os }}
 
@@ -184,7 +184,8 @@ jobs:
         run: |
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
+          OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "not unitary_correct" $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "unitary_correct" $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -203,7 +204,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: C++ tests (OpenBLAS)
     runs-on: ${{ matrix.os }}
 
@@ -259,7 +260,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: C++ tests (OpenBLAS without LAPACK)
     runs-on: ${{ matrix.os }}
 
@@ -314,7 +315,7 @@ jobs:
       matrix:
         os: [ubuntu-22.04]
         pl_backend: ["lightning_qubit"]
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: Python tests with OpenBLAS
     runs-on: ${{ matrix.os }}
 
@@ -392,7 +393,8 @@ jobs:
         run: |
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
-          PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
+          OMP_NUM_THREADS=1 PL_DEVICE=${DEVICENAME} python -m pytest -n auto tests/ -k "not unitary_correct" $COVERAGE_FLAGS
+          PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "unitary_correct" $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -419,7 +421,7 @@ jobs:
         pl_backend: ["lightning_kokkos"]
         exec_model: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.kokkos_version) }}
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: C++ tests (Kokkos)
     runs-on: ${{ matrix.os }}
 
@@ -494,7 +496,7 @@ jobs:
         exclude:
           - pl_backend: ["all"]
             exec_model: OPENMP
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: Python tests with Kokkos
     runs-on: ${{ matrix.os }}
 
@@ -609,10 +611,11 @@ jobs:
         if: ${{ matrix.pl_backend == 'all' }}
         run: |
           cd main/
-          PL_DEVICE=lightning.qubit python -m pytest tests/ $COVERAGE_FLAGS
+          OMP_NUM_THREADS=1 PL_DEVICE=lightning.qubit python -m pytest -n auto tests/ -k "not unitary_correct" $COVERAGE_FLAGS
+          PL_DEVICE=lightning.qubit python -m pytest tests/ -k "unitary_correct" $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.qubit --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
-          PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS
+          PL_DEVICE=lightning.kokkos python -m pytest tests/ $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.kokkos --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device lightning.kokkos --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}
@@ -682,7 +685,7 @@ jobs:
         os: [ubuntu-22.04]
         exec_model: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.build_and_cache_Kokkos.outputs.kokkos_version) }}
-    timeout-minutes: 30
+    timeout-minutes: 60
     name: C++ tests (multiple backends)
     runs-on: ${{ matrix.os }}
 

Makefile

@@ -6,8 +6,10 @@ PL_BACKEND ?= "$(if $(backend:-=),$(backend),lightning_qubit)"
 
 ifdef check
     CHECK := --check
+	ICHECK := --check --diff
 else
     CHECK :=
+    ICHECK :=
 endif
 
 ifdef build_options
@@ -113,7 +115,7 @@ format-cpp:
 	./bin/format $(CHECK) ./pennylane_lightning
 
 format-python:
-	isort --profile black ./pennylane_lightning/ ./mpitests ./tests $(CHECK)
+	isort --profile black ./pennylane_lightning/ ./mpitests ./tests $(ICHECK)
 	black -l 100 ./pennylane_lightning/ ./mpitests ./tests $(CHECK)
 
 .PHONY: check-tidy

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0-dev14"
+__version__ = "0.36.0-dev15"

pennylane_lightning/core/src/simulators/lightning_qubit/bindings/LQubitBindings.hpp

@@ -220,6 +220,10 @@ void registerBackendClassSpecificBindings(PyClass &pyclass) {
                 }
             },
             "Copy StateVector data into a Numpy array.")
+        .def("collapse", &StateVectorT::collapse,
+             "Collapse the statevector onto the 0 or 1 branch of a given wire.")
+        .def("normalize", &StateVectorT::normalize,
+             "Normalizes the statevector to norm 1.")
         .def("applyControlledMatrix", &applyControlledMatrix<StateVectorT>,
              "Apply controlled operation")
         .def("kernel_map", &svKernelMap<StateVectorT>,

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -64,6 +64,13 @@
     from pennylane.ops.op_math import Adjoint
     from pennylane.wires import Wires
 
+    # pylint: disable=import-error, no-name-in-module, ungrouped-imports
+    from pennylane_lightning.core._serialize import (
+        QuantumScriptSerializer,
+        global_phase_diagonal,
+    )
+    from pennylane_lightning.core._version import __version__
+
     # pylint: disable=import-error, no-name-in-module, ungrouped-imports
     from pennylane_lightning.lightning_kokkos_ops.algorithms import (
         AdjointJacobianC64,
@@ -72,13 +79,6 @@
         create_ops_listC128,
     )
 
-    # pylint: disable=import-error, no-name-in-module, ungrouped-imports
-    from pennylane_lightning.core._serialize import (
-        QuantumScriptSerializer,
-        global_phase_diagonal,
-    )
-    from pennylane_lightning.core._version import __version__
-
     def _kokkos_dtype(dtype):
         if dtype not in [np.complex128, np.complex64]:  # pragma: no cover
             raise ValueError(f"Data type is not supported for state-vector computation: {dtype}")

pennylane_lightning/lightning_qubit/_measurements.py

@@ -48,8 +48,6 @@
 
 from pennylane_lightning.core._serialize import QuantumScriptSerializer
 
-from ._state_vector import LightningStateVector
-
 
 class LightningMeasurements:
     """Lightning Measurements class
@@ -72,7 +70,7 @@ class LightningMeasurements:
 
     def __init__(
         self,
-        qubit_state: LightningStateVector,
+        qubit_state,
         mcmc: bool = None,
         kernel_name: str = None,
         num_burnin: int = None,

pennylane_lightning/lightning_qubit/_state_vector.py

@@ -29,10 +29,14 @@
 import numpy as np
 import pennylane as qml
 from pennylane import BasisState, DeviceError, StatePrep
+from pennylane.measurements import MidMeasureMP
+from pennylane.ops import Conditional
 from pennylane.ops.op_math import Adjoint
 from pennylane.tape import QuantumScript
 from pennylane.wires import Wires
 
+from ._measurements import LightningMeasurements
+
 
 class LightningStateVector:
     """Lightning state-vector class.
@@ -220,10 +224,10 @@ def _apply_lightning_controlled(self, operation):
         """Apply an arbitrary controlled operation to the state tensor.
 
         Args:
-            operation (~pennylane.operation.Operation): operation to apply
+            operation (~pennylane.operation.Operation): controlled operation to apply
 
         Returns:
-            array[complex]: the output state tensor
+            None
         """
         state = self.state_vector
 
@@ -246,14 +250,36 @@ def _apply_lightning_controlled(self, operation):
                 False,
             )
 
-    def _apply_lightning(self, operations):
+    def _apply_lightning_midmeasure(self, operation: MidMeasureMP, mid_measurements: dict):
+        """Execute a MidMeasureMP operation and return the sample in mid_measurements.
+        Args:
+            operation (~pennylane.operation.Operation): mid-circuit measurement
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
+        Returns:
+            None
+        """
+        wires = self.wires.indices(operation.wires)
+        wire = list(wires)[0]
+        circuit = QuantumScript([], [qml.sample(wires=operation.wires)], shots=1)
+        sample = LightningMeasurements(self).measure_final_state(circuit)
+        sample = np.squeeze(sample)
+        if operation.postselect is not None and sample != operation.postselect:
+            mid_measurements[operation] = -1
+            return
+        mid_measurements[operation] = sample
+        getattr(self.state_vector, "collapse")(wire, bool(sample))
+        if operation.reset and bool(sample):
+            self.apply_operations([qml.PauliX(operation.wires)], mid_measurements=mid_measurements)
+
+    def _apply_lightning(self, operations, mid_measurements: dict = None):
         """Apply a list of operations to the state tensor.
 
         Args:
             operations (list[~pennylane.operation.Operation]): operations to apply
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
 
         Returns:
-            array[complex]: the output state tensor
+            None
         """
         state = self.state_vector
 
@@ -271,7 +297,12 @@ def _apply_lightning(self, operations):
             method = getattr(state, name, None)
             wires = list(operation.wires)
 
-            if method is not None:  # apply specialized gate
+            if isinstance(operation, Conditional):
+                if operation.meas_val.concretize(mid_measurements):
+                    self._apply_lightning([operation.then_op])
+            elif isinstance(operation, MidMeasureMP):
+                self._apply_lightning_midmeasure(operation, mid_measurements)
+            elif method is not None:  # apply specialized gate
                 param = operation.parameters
                 method(wires, invert_param, param)
             elif isinstance(operation, qml.ops.Controlled):  # apply n-controlled gate
@@ -286,7 +317,7 @@ def _apply_lightning(self, operations):
                     # To support older versions of PL
                     method(operation.matrix, wires, False)
 
-    def apply_operations(self, operations):
+    def apply_operations(self, operations, mid_measurements: dict = None):
         """Applies operations to the state vector."""
         # State preparation is currently done in Python
         if operations:  # make sure operations[0] exists
@@ -297,21 +328,22 @@ def apply_operations(self, operations):
                 self._apply_basis_state(operations[0].parameters[0], operations[0].wires)
                 operations = operations[1:]
 
-        self._apply_lightning(operations)
+        self._apply_lightning(operations, mid_measurements=mid_measurements)
 
-    def get_final_state(self, circuit: QuantumScript):
+    def get_final_state(self, circuit: QuantumScript, mid_measurements: dict = None):
         """
         Get the final state that results from executing the given quantum script.
 
         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:
             LightningStateVector: Lightning final state class.
 
         """
-        self.apply_operations(circuit.operations)
+        self.apply_operations(circuit.operations, mid_measurements=mid_measurements)
 
         return self

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -25,12 +25,14 @@
 from pennylane.devices.modifiers import simulator_tracking, single_tape_support
 from pennylane.devices.preprocess import (
     decompose,
+    mid_circuit_measurements,
     no_sampling,
     validate_adjoint_trainable_params,
     validate_device_wires,
     validate_measurements,
     validate_observables,
 )
+from pennylane.measurements import MidMeasureMP
 from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.transforms.core import TransformProgram
 from pennylane.typing import Result, ResultBatch
@@ -71,6 +73,16 @@ def simulate(circuit: QuantumScript, state: LightningStateVector, mcmc: dict = N
     if mcmc is None:
         mcmc = {}
     state.reset_state()
+    has_mcm = any(isinstance(op, MidMeasureMP) for op in circuit.operations)
+    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,
+        )
     final_state = state.get_final_state(circuit)
     return LightningMeasurements(final_state, **mcmc).measure_final_state(circuit)
 
@@ -200,6 +212,8 @@ def simulate_and_jacobian(circuit: QuantumTape, state: LightningStateVector, bat
         "QFT",
         "ECR",
         "BlockEncode",
+        "MidMeasureMP",
+        "Conditional",
     }
 )
 # The set of supported operations.
@@ -432,7 +446,7 @@ def preprocess(self, execution_config: ExecutionConfig = DefaultExecutionConfig)
         program.add_transform(validate_measurements, name=self.name)
         program.add_transform(validate_observables, accepted_observables, name=self.name)
         program.add_transform(validate_device_wires, self.wires, name=self.name)
-        program.add_transform(qml.defer_measurements, device=self)
+        program.add_transform(mid_circuit_measurements, device=self)
         program.add_transform(decompose, stopping_condition=stopping_condition, name=self.name)
         program.add_transform(qml.transforms.broadcast_expand)
 

tests/new_api/test_device.py

@@ -233,7 +233,9 @@ def test_preprocess(self, adjoint):
         expected_program.add_transform(validate_measurements, name=device.name)
         expected_program.add_transform(validate_observables, accepted_observables, name=device.name)
         expected_program.add_transform(validate_device_wires, device.wires, name=device.name)
-        expected_program.add_transform(qml.defer_measurements, device=device)
+        expected_program.add_transform(
+            qml.devices.preprocess.mid_circuit_measurements, device=device
+        )
         expected_program.add_transform(
             decompose, stopping_condition=stopping_condition, name=device.name
         )