Add test_templates module + Fix LQ decomposition of StatePrep Ops + Fix LQ decomposition strategy of QFT and GroverOperator (#684)

* Add test for QSVT.

* Auto update version

* Update changelog.

* Add tests for embedding, layer and stateprep templates.

* Parametrize tests over n_qubits.

* Add qchem template tests.

* Add a few misc template tests.

* Auto update version

* trigger ci

* Fix serialize.

* Fix formatting.

* Update tests/test_templates.py

Co-authored-by: Amintor Dusko <[email protected]>

* Update tests/test_templates.py

Co-authored-by: Amintor Dusko <[email protected]>

* Add xfail condition for AmplitudeEmbedding.

* Fix pytest.skip top cond.

* Update tests/test_templates.py

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

* Update pennylane_lightning/core/_serialize.py

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

* Add breaking tests for if not observable

* trigger ci

* Auto update version

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

* Fix test params

---------

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

Author: 5 people

.github/CHANGELOG.md

@@ -52,6 +52,9 @@
 
 ### Improvements
 
+* Add `test_templates.py` module where Grover and QSVT are tested.
+  [(#684)](https://github.com/PennyLaneAI/pennylane-lightning/pull/684)
+
 * Create `cuda_utils` for common usage of CUDA related backends.
   [(#676)](https://github.com/PennyLaneAI/pennylane-lightning/pull/676)
 
@@ -74,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/_serialize.py

@@ -255,6 +255,9 @@ def map_wire(wire: int):
 
     def _pauli_sentence(self, observable, wires_map: dict = None):
         """Serialize a :class:`pennylane.pauli.PauliSentence` into a Hamiltonian."""
+        # Trivial Pauli sentences' items is empty, cannot unpack
+        if not observable:
+            return self.hamiltonian_obs(np.array([0.0]).astype(self.rtype), [self._ob(Identity(0))])
         pwords, coeffs = zip(*observable.items())
         terms = [self._pauli_word(pw, wires_map) for pw in pwords]
         coeffs = np.array(coeffs).astype(self.rtype)

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_adjoint_jacobian.py

@@ -1481,6 +1481,13 @@ def convert_to_array_gpu(params):
     def convert_to_array_gpu_default(params):
         return np.hstack(qnode_gpu_default(params))
 
+    i_cpu = qnode_cpu(params)
+    i_gpu = qnode_gpu(params)
+    i_gpu_default = qnode_gpu_default(params)
+
+    assert np.allclose(i_cpu, i_gpu)
+    assert np.allclose(i_gpu, i_gpu_default)
+
     j_cpu = qml.jacobian(qnode_cpu)(params)
     j_gpu = qml.jacobian(qnode_gpu)(params)
     j_gpu_default = qml.jacobian(qnode_gpu_default)(params)

tests/test_execute.py

@@ -87,40 +87,3 @@ def dev_l_execute(t):
 
         assert np.allclose(grad_dev_l, grad_qml_l, tol)
         assert np.allclose(grad_dev_l, grad_qml_d, tol)
-
-
-class TestGrover:
-    """Test Grover's algorithm (multi-controlled gates, decomposition, etc.)"""
-
-    @pytest.mark.parametrize("num_qubits", range(4, 8))
-    def test_grover(self, num_qubits):
-        np.random.seed(42)
-        omega = np.random.rand(num_qubits) > 0.5
-        dev = qml.device(device_name, wires=num_qubits)
-        wires = list(range(num_qubits))
-
-        @qml.qnode(dev, diff_method=None)
-        def circuit(omega):
-            iterations = int(np.round(np.sqrt(2**num_qubits) * np.pi / 4))
-
-            # Initial state preparation
-            for wire in wires:
-                qml.Hadamard(wires=wire)
-
-            # Grover's iterator
-            for _ in range(iterations):
-                qml.FlipSign(omega, wires=wires)
-                qml.templates.GroverOperator(wires)
-
-            return qml.probs(wires=wires)
-
-        prob = circuit(omega)
-        index = omega.astype(int)
-        index = functools.reduce(
-            lambda sum, x: sum + (1 << x[0]) * x[1],
-            zip([i for i in range(len(index) - 1, -1, -1)], index),
-            0,
-        )
-        assert np.allclose(np.sum(prob), 1.0)
-        assert prob[index] > 0.95
-        assert np.sum(prob) - prob[index] < 0.05