Fix decompositions with LightningQubit (#687)

* Fix bug; add tests; update changelog

* Auto update version

* Update tests

* Update qft/grover decomp

* Addressing code review

* Update pennylane_lightning/lightning_qubit/lightning_qubit.py

Co-authored-by: Ali Asadi <[email protected]>

* Trigger CI

* Auto update version

---------

Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Ali Asadi <[email protected]>

Author: 3 people

.github/CHANGELOG.md

@@ -77,6 +77,12 @@
 
 ### Bug fixes
 
+* `LightningQubit` correctly decomposes state prep operations when used in the middle of a circuit.
+  [(#687)](https://github.com/PennyLaneAI/pennylane/pull/687)
+
+* `LightningQubit` correctly decomposes `qml.QFT` and `qml.GroverOperator` if `len(wires)` is greater than 9 and 12 respectively.
+  [(#687)](https://github.com/PennyLaneAI/pennylane/pull/687)
+
 * Specify `isort` `--py` (Python version) and `-l` (max line length) to stabilize `isort` across Python versions and environments.
   [(#647)](https://github.com/PennyLaneAI/pennylane-lightning/pull/647)
 

pennylane_lightning/core/_version.py

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

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -182,9 +182,6 @@ def simulate_and_vjp(
 _operations = frozenset(
     {
         "Identity",
-        "BasisState",
-        "QubitStateVector",
-        "StatePrep",
         "QubitUnitary",
         "ControlledQubitUnitary",
         "MultiControlledX",
@@ -292,13 +289,20 @@ def simulate_and_vjp(
 
 def stopping_condition(op: Operator) -> bool:
     """A function that determines whether or not an operation is supported by ``lightning.qubit``."""
+    # These thresholds are adapted from `lightning_base.py`
+    # To avoid building matrices beyond the given thresholds.
+    # This should reduce runtime overheads for larger systems.
+    if isinstance(op, qml.QFT):
+        return len(op.wires) < 10
+    if isinstance(op, qml.GroverOperator):
+        return len(op.wires) < 13
     return op.name in _operations
 
 
 def stopping_condition_shots(op: Operator) -> bool:
     """A function that determines whether or not an operation is supported by ``lightning.qubit``
     with finite shots."""
-    return op.name in _operations or isinstance(op, (MidMeasureMP, qml.ops.op_math.Conditional))
+    return stopping_condition(op) or isinstance(op, (MidMeasureMP, qml.ops.op_math.Conditional))
 
 
 def accepted_observables(obs: Operator) -> bool:
@@ -536,6 +540,7 @@ def preprocess(self, execution_config: ExecutionConfig = DefaultExecutionConfig)
             decompose,
             stopping_condition=stopping_condition,
             stopping_condition_shots=stopping_condition_shots,
+            skip_initial_state_prep=True,
             name=self.name,
         )
         program.add_transform(qml.transforms.broadcast_expand)

tests/new_api/test_device.py

@@ -31,6 +31,7 @@
     adjoint_measurements,
     adjoint_observables,
     decompose,
+    mid_circuit_measurements,
     no_sampling,
     stopping_condition,
     stopping_condition_shots,
@@ -258,13 +259,12 @@ 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.devices.preprocess.mid_circuit_measurements, device=device
-        )
+        expected_program.add_transform(mid_circuit_measurements, device=device)
         expected_program.add_transform(
             decompose,
             stopping_condition=stopping_condition,
             stopping_condition_shots=stopping_condition_shots,
+            skip_initial_state_prep=True,
             name=device.name,
         )
         expected_program.add_transform(qml.transforms.broadcast_expand)
@@ -293,6 +293,63 @@ def test_preprocess(self, adjoint):
         actual_program, _ = device.preprocess(config)
         assert actual_program == expected_program
 
+    @pytest.mark.parametrize(
+        "op, is_trainable",
+        [
+            (qml.StatePrep([1 / np.sqrt(2), 1 / np.sqrt(2)], wires=0), False),
+            (qml.StatePrep(qml.numpy.array([1 / np.sqrt(2), 1 / np.sqrt(2)]), wires=0), True),
+            (qml.StatePrep(np.array([1, 0]), wires=0), False),
+            (qml.BasisState([1, 1], wires=[0, 1]), False),
+            (qml.BasisState(qml.numpy.array([1, 1]), wires=[0, 1]), True),
+        ],
+    )
+    def test_preprocess_state_prep_first_op_decomposition(self, op, is_trainable):
+        """Test that state prep ops in the beginning of a tape are decomposed with adjoint
+        but not otherwise."""
+        tape = qml.tape.QuantumScript([op, qml.RX(1.23, wires=0)], [qml.expval(qml.PauliZ(0))])
+        device = LightningDevice(wires=3)
+
+        if is_trainable:
+            # Need to decompose twice as the state prep ops we use first decompose into a template
+            decomp = op.decomposition()[0].decomposition()
+        else:
+            decomp = [op]
+
+        config = ExecutionConfig(gradient_method="best" if is_trainable else None)
+        program, _ = device.preprocess(config)
+        [new_tape], _ = program([tape])
+        expected_tape = qml.tape.QuantumScript([*decomp, qml.RX(1.23, wires=0)], tape.measurements)
+        assert qml.equal(new_tape, expected_tape)
+
+    @pytest.mark.parametrize(
+        "op, decomp_depth",
+        [
+            (qml.StatePrep([1 / np.sqrt(2), 1 / np.sqrt(2)], wires=0), 1),
+            (qml.StatePrep(np.array([1, 0]), wires=0), 1),
+            (qml.BasisState([1, 1], wires=[0, 1]), 1),
+            (qml.BasisState(qml.numpy.array([1, 1]), wires=[0, 1]), 1),
+            (qml.AmplitudeEmbedding([1 / np.sqrt(2), 1 / np.sqrt(2)], wires=0), 2),
+            (qml.MottonenStatePreparation([1 / np.sqrt(2), 1 / np.sqrt(2)], wires=0), 0),
+        ],
+    )
+    def test_preprocess_state_prep_middle_op_decomposition(self, op, decomp_depth):
+        """Test that state prep ops in the middle of a tape are always decomposed."""
+        tape = qml.tape.QuantumScript(
+            [qml.RX(1.23, wires=0), op, qml.CNOT([0, 1])], [qml.expval(qml.PauliZ(0))]
+        )
+        device = LightningDevice(wires=3)
+
+        for _ in range(decomp_depth):
+            op = op.decomposition()[0]
+        decomp = op.decomposition()
+
+        program, _ = device.preprocess()
+        [new_tape], _ = program([tape])
+        expected_tape = qml.tape.QuantumScript(
+            [qml.RX(1.23, wires=0), *decomp, qml.CNOT([0, 1])], tape.measurements
+        )
+        assert qml.equal(new_tape, expected_tape)
+
     @pytest.mark.usefixtures("use_legacy_and_new_opmath")
     @pytest.mark.parametrize("theta, phi", list(zip(THETA, PHI)))
     @pytest.mark.parametrize(

tests/test_templates.py

@@ -62,6 +62,23 @@ def circuit(omega):
         assert np.allclose(np.sum(prob), 1.0)
         assert prob[index] > 0.95
 
+    @pytest.mark.skipif(not LightningDevice._new_API, reason="New API required.")
+    @pytest.mark.parametrize("wires", [5, 10, 13, 15])
+    def test_preprocess_grover_operator_decomposition(self, wires):
+        """Test that qml.GroverOperator is not decomposed for less than 10 wires."""
+        tape = qml.tape.QuantumScript(
+            [qml.GroverOperator(wires=list(range(wires)))], [qml.expval(qml.PauliZ(0))]
+        )
+        dev = LightningDevice(wires=wires)
+
+        program, _ = dev.preprocess()
+        [new_tape], _ = program([tape])
+
+        if wires >= 13:
+            assert all(not isinstance(op, qml.GroverOperator) for op in new_tape.operations)
+        else:
+            assert tape.operations == [qml.GroverOperator(wires=list(range(wires)))]
+
 
 class TestAngleEmbedding:
     """Test the AngleEmbedding algorithm."""
@@ -416,7 +433,6 @@ class TestGateFabric:
     """Test the GateFabric algorithm."""
 
     def test_gatefabric(self):
-
         # Build the electronic Hamiltonian
         symbols = ["H", "H"]
         coordinates = np.array([0.0, 0.0, -0.6614, 0.0, 0.0, 0.6614])
@@ -446,7 +462,6 @@ class TestUCCSD:
     """Test the UCCSD algorithm."""
 
     def test_uccsd(self):
-
         # Define the molecule
         symbols = ["H", "H", "H"]
         geometry = np.array(
@@ -490,7 +505,6 @@ class TestkUpCCGSD:
     """Test the kUpCCGSD algorithm."""
 
     def test_kupccgsd(self):
-
         # Define the molecule
         symbols = ["H", "H", "H"]
         geometry = np.array(
@@ -533,7 +547,6 @@ class TestParticleConservingU1:
     """Test the ParticleConservingU1 algorithm."""
 
     def test_particleconservingu1(self):
-
         # Build the electronic Hamiltonian
         symbols, coordinates = (["H", "H"], np.array([0.0, 0.0, -0.66140414, 0.0, 0.0, 0.66140414]))
         _, n_qubits = qml.qchem.molecular_hamiltonian(symbols, coordinates)
@@ -567,7 +580,6 @@ class TestParticleConservingU2:
     """Test the ParticleConservingU2 algorithm."""
 
     def test_particleconservingu2(self):
-
         # Build the electronic Hamiltonian
         symbols, coordinates = (["H", "H"], np.array([0.0, 0.0, -0.66140414, 0.0, 0.0, 0.66140414]))
         _, n_qubits = qml.qchem.molecular_hamiltonian(symbols, coordinates)
@@ -668,7 +680,6 @@ class TestQuantumPhaseEstimation:
 
     @pytest.mark.parametrize("n_qubits", range(2, 14, 2))
     def test_quantumphaseestimation(self, n_qubits):
-
         phase = 5
         target_wires = [0]
         unitary = qml.RX(phase, wires=0).matrix()
@@ -701,7 +712,6 @@ class TestQFT:
 
     @pytest.mark.parametrize("n_qubits", range(2, 14, 2))
     def test_qft(self, n_qubits):
-
         dev = qml.device(device_name, wires=n_qubits)
         dq = qml.device("default.qubit")
 
@@ -717,13 +727,29 @@ def circuit(basis_state):
 
         assert np.allclose(res, ref)
 
+    @pytest.mark.skipif(not LightningDevice._new_API, reason="New API required")
+    @pytest.mark.parametrize("wires", [5, 9, 10, 13])
+    def test_preprocess_qft_decomposition(self, wires):
+        """Test that qml.QFT is not decomposed for less than 10 wires."""
+        tape = qml.tape.QuantumScript(
+            [qml.QFT(wires=list(range(wires)))], [qml.expval(qml.PauliZ(0))]
+        )
+        dev = LightningDevice(wires=wires)
+
+        program, _ = dev.preprocess()
+        [new_tape], _ = program([tape])
+
+        if wires >= 10:
+            assert all(not isinstance(op, qml.QFT) for op in new_tape.operations)
+        else:
+            assert tape.operations == [qml.QFT(wires=list(range(wires)))]
+
 
 class TestAQFT:
     """Test the AQFT algorithm."""
 
     @pytest.mark.parametrize("n_qubits", range(4, 14, 2))
     def test_aqft(self, n_qubits):
-
         dev = qml.device(device_name, wires=n_qubits)
         dq = qml.device("default.qubit")