Fix decompositions with LightningQubit

.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 wires are above a certain threshold.
+ [(#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/lightning_qubit/lightning_qubit.py

@@ -182,9 +182,6 @@
 _operations = frozenset(
  {
  "Identity",
- "BasisState",
- "QubitStateVector",
- "StatePrep",
  "QubitUnitary",
  "ControlledQubitUnitary",
  "MultiControlledX",
@@ -292,6 +289,11 @@
 
 def stopping_condition(op: Operator) -> bool:
  """A function that determines whether or not an operation is supported by ``lightning.qubit``."""
+ 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
 
 
@@ -536,6 +538,7 @@
  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

@@ -249,6 +249,95 @@ def test_preprocess_correct_config_setup(self, config, expected_config):
 
  assert new_config == expected_config
 
+ @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.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)))]
+
+ @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)))]
+
  @pytest.mark.parametrize("adjoint", [True, False])
  def test_preprocess(self, adjoint):
  """Test that the transform program returned by preprocess is correct"""