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

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