Add StatePrep support (#137)

* Add  support

* requirements for ci

* change log

* fix requirements bug?

* Update requirements.txt

* ci

---------

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>

CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Improvements 🛠
 
+* Added support for `qml.StatePrep` as a state preparation operation.
+  [(#137)](https://github.com/PennyLaneAI/pennylane-rigetti/pull/137)
+
 ### Breaking changes 💔
 
 ### Deprecations 👋
@@ -16,6 +19,8 @@
 
 This release contains contributions from (in alphabetical order):
 
+ Jay Soni
+
 ---
 
 # Release 0.31.0

doc/devices/numpy_wavefunction.rst

@@ -58,4 +58,4 @@ array([0.97517033, 0.04904283])
 Supported operations
 ~~~~~~~~~~~~~~~~~~~~
 
-All Rigetti devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``QubitStateVector`` state preparation operation.
+All Rigetti devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``StatePrepBase`` state preparation operations.

doc/devices/qpu.rst

@@ -48,7 +48,7 @@ We can then integrate the quantum hardware and PennyLane's automatic differentia
 Supported operations
 ~~~~~~~~~~~~~~~~~~~~
 
-All devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``QubitStateVector`` state preparation operation.
+All devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``StatePrepBase`` state preparation operations.
 
 quilc server configuration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~

doc/devices/qvm.rst

@@ -96,7 +96,7 @@ Printing out the results of the three device expectation values:
 Supported operations
 ~~~~~~~~~~~~~~~~~~~~
 
-All devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``QubitStateVector`` state preparation operation.
+All devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``StatePrepBase`` state preparation operations.
 
 Supported observables
 ~~~~~~~~~~~~~~~~~~~~~

doc/devices/wavefunction.rst

@@ -53,6 +53,6 @@ array([0.97517033, 0.04904283])
 Supported operations
 ~~~~~~~~~~~~~~~~~~~~
 
-All devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``QubitStateVector`` state preparation operation.
+All devices support all PennyLane `operations and observables <https://pennylane.readthedocs.io/en/stable/introduction/operations.html#qubit-operations>`_, with the exception of the PennyLane ``StatePrepBase`` state preparation operations.
 
 .. include:: ./qvm_and_quilc_server_configuration.rst

pennylane_rigetti/device.py

@@ -253,7 +253,7 @@ def apply_circuit_operations(self, operations):
                     # Array not supported
                     par = [float(i) for i in par]
 
-            if i > 0 and operation.name in ("QubitStateVector", "BasisState"):
+            if i > 0 and operation.name in ("QubitStateVector", "StatePrep", "BasisState"):
                 name = operation.name
                 short_name = self.short_name
                 raise DeviceError(

pennylane_rigetti/qc.py

@@ -215,7 +215,7 @@ def apply_parametric_operations(self, operations):
             # map the operation wires to the physical device qubits
             device_wires = self.map_wires(operation.wires)
 
-            if i > 0 and operation.name in ("QubitStateVector", "BasisState"):
+            if i > 0 and operation.name in ("QubitStateVector", "StatePrep", "BasisState"):
                 raise DeviceError(
                     f"Operation {operation.name} cannot be used after other Operations have already been applied "
                     f"on a {self.short_name} device."

requirements.txt

@@ -1,5 +1,5 @@
+pennylane @ git+https://github.com/PennyLaneAI/pennylane.git
 numpy>=1.21,<1.24
 networkx>=2.5,<3.0
-pennylane>=0.18
 pyquil>=3.0.0,<4.0.0
 qcs-api-client>=0.20.13,<0.22.0

tests/test_device.py

@@ -125,13 +125,14 @@ def test_apply_basis_state_raises_an_error_if_not_first(self):
         ):
             dev.apply(queue)
 
-    def test_apply_qubitstatesvector_raises_an_error_if_not_first(self):
-        """Test that there is an error raised when the QubitStateVector is not
+    @pytest.mark.parametrize("op", [qml.StatePrep, qml.QubitStateVector])
+    def test_apply_statesvector_raises_an_error_if_not_first(self, op):
+        """Test that there is an error raised when the QubitStateVector or StatePrep op is not
         applied as the first operation."""
         wires = 1
         dev = RigettiDevice(wires=wires, shots=1)
 
-        operation = qml.QubitStateVector(np.array([1, 0]), wires=list(range(wires)))
+        operation = op(np.array([1, 0]), wires=list(range(wires)))
         queue = [qml.PauliX(0), operation]
         with pytest.raises(
             qml.DeviceError,

tests/test_qvm.py

@@ -606,12 +606,13 @@ def test_parametric_compilation_with_numeric_and_symbolic_queue(
         assert len(dev._compiled_program_dict.items()) == 1
         assert mock_qvm.compile.call_count == 1
 
-    def test_apply_qubitstatesvector_raises_an_error_if_not_first(self):
-        """Test that there is an error raised when the QubitStateVector is not
+    @pytest.mark.parametrize("op", [qml.QubitStateVector, qml.StatePrep])
+    def test_apply_qubitstatesvector_raises_an_error_if_not_first(self, op):
+        """Test that there is an error raised when the QubitStateVector or StatPrep is not
         applied as the first operation."""
         dev = qml.device("rigetti.qvm", device="2q-qvm", parametric_compilation=True)
 
-        operation = qml.QubitStateVector(np.array([1, 0]), wires=[0])
+        operation = op(np.array([1, 0]), wires=[0])
         queue = [qml.PauliX(0), operation]
         with pytest.raises(
             qml.DeviceError,