QubitStateVector => StatePrep + parametrize tests.

.github/CHANGELOG.md

@@ -16,6 +16,9 @@
 
 ### Breaking changes
 
+* Rename `QubitStateVector` to `StatePrep` in the `LightningQubit` and `LightningKokkos` classes.
+  [(#486)](https://github.com/PennyLaneAI/pennylane-lightning/pull/486)
+
 * Modify `adjointJacobian` methods to accept a (maybe unused) reference `StateVectorT`, allowing device-backed simulators to directly access state vector data for adjoint differentiation instead of copying it back-and-forth into `JacobianData` (host memory).
   [(#477)](https://github.com/PennyLaneAI/pennylane-lightning/pull/477)
 

.readthedocs.yml

@@ -7,7 +7,7 @@ python:
   install:
     - requirements: ci_build_requirements.txt
     - requirements: doc/requirements.txt
-    - requirements: requirements.txt
+    - requirements: requirements-dev.txt
     - method: pip
       path: .
   system_packages: true

pennylane_lightning/core/_serialize.py

@@ -23,7 +23,7 @@
     PauliY,
     PauliZ,
     Identity,
-    QubitStateVector,
+    StatePrep,
     Rot,
 )
 from pennylane.operation import Tensor
@@ -206,7 +206,7 @@ def serialize_ops(
         uses_stateprep = False
 
         for operation in tape.operations:
-            if isinstance(operation, (BasisState, QubitStateVector)):
+            if isinstance(operation, (BasisState, StatePrep)):
                 uses_stateprep = True
                 continue
             if isinstance(operation, Rot):

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.32.0-dev13"
+__version__ = "0.32.0-dev14"

pennylane_lightning/core/lightning_base.py

@@ -25,7 +25,7 @@
 from pennylane import (
     BasisState,
     QubitDevice,
-    QubitStateVector,
+    StatePrep,
 )
 from pennylane.devices import DefaultQubit
 from pennylane.measurements import MeasurementProcess
@@ -277,7 +277,7 @@ def _process_jacobian_tape(self, tape, starting_state, use_device_state):
             # get op_idx-th operator among differentiable operators
             operation, _, _ = tape.get_operation(op_idx)
             if isinstance(operation, Operation) and not isinstance(
-                operation, (BasisState, QubitStateVector)
+                operation, (BasisState, StatePrep)
             ):
                 # We now just ignore non-op or state preps
                 tp_shift.append(trainable_param)

pennylane_lightning/core/src/algorithms/JacobianData.hpp

@@ -246,7 +246,7 @@ template <class StateVectorT> class JacobianData {
      * we want to take a derivative respect to the :math:`i`-th operation.
      *
      * Further note that ``ops`` does not contain state preparation operations
-     * (e.g. QubitStateVector) or Hamiltonian coefficients.
+     * (e.g. StatePrep) or Hamiltonian coefficients.
      * @endrst
      */
     JacobianData(size_t num_params, size_t num_elem, const ComplexT *sv_ptr,

pennylane_lightning/core/src/simulators/lightning_kokkos/algorithms/AdjointJacobianKokkos.hpp

@@ -134,6 +134,7 @@ class AdjointJacobian final
                         "The operation is not supported using the adjoint "
                         "differentiation method");
             if ((ops_name[op_idx] == "QubitStateVector") ||
+                (ops_name[op_idx] == "StatePrep") ||
                 (ops_name[op_idx] == "BasisState")) {
                 continue;
             }

pennylane_lightning/core/src/simulators/lightning_qubit/algorithms/AdjointJacobianLQubit.hpp

@@ -285,6 +285,7 @@ class AdjointJacobian final
                         "The operation is not supported using the adjoint "
                         "differentiation method");
             if ((ops_name[op_idx] == "QubitStateVector") ||
+                (ops_name[op_idx] == "StatePrep") ||
                 (ops_name[op_idx] == "BasisState")) {
                 continue; // Ignore them
             }

pennylane_lightning/core/src/simulators/lightning_qubit/algorithms/VectorJacobianProduct.hpp

@@ -121,6 +121,7 @@ class VectorJacobianProduct final
                         "The operation is not supported using the adjoint "
                         "differentiation method");
             if ((ops_name[op_idx] == "QubitStateVector") ||
+                (ops_name[op_idx] == "StatePrep") ||
                 (ops_name[op_idx] == "BasisState")) {
                 continue; // ignore them
             }

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -52,7 +52,7 @@
     from pennylane import (
         math,
         BasisState,
-        QubitStateVector,
+        StatePrep,
         Projector,
         Rot,
         DeviceError,
@@ -89,6 +89,7 @@ def _kokkos_configuration():
         "Identity",
         "BasisState",
         "QubitStateVector",
+        "StatePrep",
         "QubitUnitary",
         "ControlledQubitUnitary",
         "MultiControlledX",
@@ -429,7 +430,7 @@ def apply_lightning(self, operations):
         def apply(self, operations, rotations=None, **kwargs):
             # State preparation is currently done in Python
             if operations:  # make sure operations[0] exists
-                if isinstance(operations[0], QubitStateVector):
+                if isinstance(operations[0], StatePrep):
                     self._apply_state_vector(
                         operations[0].parameters[0].copy(), operations[0].wires
                     )
@@ -439,7 +440,7 @@ def apply(self, operations, rotations=None, **kwargs):
                     operations = operations[1:]
 
             for operation in operations:
-                if isinstance(operation, (QubitStateVector, BasisState)):
+                if isinstance(operation, (StatePrep, BasisState)):
                     raise DeviceError(
                         f"Operation {operation.name} cannot be used after other "
                         + f"Operations have already been applied on a {self.short_name} device."

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -50,7 +50,7 @@
     from pennylane import (
         math,
         BasisState,
-        QubitStateVector,
+        StatePrep,
         Projector,
         Rot,
         DeviceError,
@@ -80,6 +80,7 @@
         "Identity",
         "BasisState",
         "QubitStateVector",
+        "StatePrep",
         "QubitUnitary",
         "ControlledQubitUnitary",
         "MultiControlledX",
@@ -372,7 +373,7 @@ def apply_lightning(self, state, operations):
         def apply(self, operations, rotations=None, **kwargs):
             # State preparation is currently done in Python
             if operations:  # make sure operations[0] exists
-                if isinstance(operations[0], QubitStateVector):
+                if isinstance(operations[0], StatePrep):
                     self._apply_state_vector(
                         operations[0].parameters[0].copy(), operations[0].wires
                     )
@@ -382,7 +383,7 @@ def apply(self, operations, rotations=None, **kwargs):
                     operations = operations[1:]
 
             for operation in operations:
-                if isinstance(operation, (QubitStateVector, BasisState)):
+                if isinstance(operation, (StatePrep, BasisState)):
                     raise DeviceError(
                         f"Operation {operation.name} cannot be used after other "
                         f"Operations have already been applied on a {self.short_name} device."

tests/test_adjoint_jacobian.py

@@ -143,15 +143,16 @@ def test_proj_unsupported(self, dev):
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
     @pytest.mark.parametrize("G", [qml.RX, qml.RY, qml.RZ])
-    def test_pauli_rotation_gradient(self, G, theta, dev):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_pauli_rotation_gradient(self, stateprep, G, theta, dev):
         """Tests that the automatic gradients of Pauli rotations are correct."""
 
         random_state = np.array(
             [0.43593284 - 0.02945156j, 0.40812291 + 0.80158023j], requires_grad=False
         )
 
         tape = qml.tape.QuantumScript(
-            [G(theta, 0)], [qml.expval(qml.PauliZ(0))], [qml.QubitStateVector(random_state, 0)]
+            [G(theta, 0)], [qml.expval(qml.PauliZ(0))], [stateprep(random_state, 0)]
         )
 
         tape.trainable_params = {1}
@@ -166,14 +167,15 @@ def test_pauli_rotation_gradient(self, G, theta, dev):
         assert np.allclose(calculated_val, numeric_val, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
-    def test_Rot_gradient(self, theta, dev):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_Rot_gradient(self, stateprep, theta, dev):
         """Tests that the device gradient of an arbitrary Euler-angle-parameterized gate is
         correct."""
 
         params = np.array([theta, theta**3, np.sqrt(2) * theta])
 
         with qml.tape.QuantumTape() as tape:
-            qml.QubitStateVector(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
+            stateprep(np.array([1.0, -1.0], requires_grad=False) / np.sqrt(2), wires=0)
             qml.Rot(*params, wires=[0])
             qml.expval(qml.PauliZ(0))
 

tests/test_apply.py

@@ -177,19 +177,19 @@ def test_apply_operation_preserve_pointer_three_wires_no_parameters(
     @pytest.mark.parametrize(
         "operation,expected_output,par",
         [
-            (qml.BasisState, [0, 0, 1, 0], [1, 0]),
             (qml.BasisState, [0, 0, 1, 0], [1, 0]),
             (qml.BasisState, [0, 0, 0, 1], [1, 1]),
             (qml.QubitStateVector, [0, 0, 1, 0], [0, 0, 1, 0]),
-            (qml.QubitStateVector, [0, 0, 1, 0], [0, 0, 1, 0]),
             (qml.QubitStateVector, [0, 0, 0, 1], [0, 0, 0, 1]),
+            (qml.StatePrep, [0, 0, 1, 0], [0, 0, 1, 0]),
+            (qml.StatePrep, [0, 0, 0, 1], [0, 0, 0, 1]),
             (
-                qml.QubitStateVector,
+                qml.StatePrep,
                 [1 / math.sqrt(3), 0, 1 / math.sqrt(3), 1 / math.sqrt(3)],
                 [1 / math.sqrt(3), 0, 1 / math.sqrt(3), 1 / math.sqrt(3)],
             ),
             (
-                qml.QubitStateVector,
+                qml.StatePrep,
                 [1 / math.sqrt(3), 0, -1 / math.sqrt(3), 1 / math.sqrt(3)],
                 [1 / math.sqrt(3), 0, -1 / math.sqrt(3), 1 / math.sqrt(3)],
             ),
@@ -454,20 +454,19 @@ def test_apply_operation_preserve_pointer_two_wires_with_parameters(
 
         assert pointer_before == pointer_after
 
-    def test_apply_errors_qubit_state_vector(self, qubit_device):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_apply_errors_qubit_state_vector(self, stateprep, qubit_device):
         """Test that apply fails for incorrect state preparation, and > 2 qubit gates"""
         dev = qubit_device(wires=2)
         with pytest.raises(ValueError, match="Sum of amplitudes-squared does not equal one."):
-            dev.apply([qml.QubitStateVector(np.array([1, -1]), wires=[0])])
+            dev.apply([stateprep(np.array([1, -1]), wires=[0])])
 
         with pytest.raises(
             DeviceError,
-            match="Operation QubitStateVector cannot be used after other Operations have already been applied ",
+            match=f"Operation {stateprep(np.array([1, 0]), wires=[0]).name} cannot be used after other Operations have already been applied ",
         ):
             dev.reset()
-            dev.apply(
-                [qml.RZ(0.5, wires=[0]), qml.QubitStateVector(np.array([0, 1, 0, 0]), wires=[0, 1])]
-            )
+            dev.apply([qml.RZ(0.5, wires=[0]), stateprep(np.array([0, 1, 0, 0]), wires=[0, 1])])
 
     def test_apply_errors_basis_state(self, qubit_device):
         dev = qubit_device(wires=2)
@@ -512,15 +511,16 @@ class TestExpval:
             (qml.Identity, [1 / math.sqrt(2), -1 / math.sqrt(2)], 1),
         ],
     )
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
     def test_expval_single_wire_no_parameters(
-        self, qubit_device, tol, operation, input, expected_output
+        self, qubit_device, tol, stateprep, operation, input, expected_output
     ):
         """Tests that expectation values are properly calculated for single-wire observables without parameters."""
         dev = qubit_device(wires=1)
         obs = operation(wires=[0])
 
         dev.reset()
-        dev.apply([qml.QubitStateVector(np.array(input), wires=[0])], obs.diagonalizing_gates())
+        dev.apply([stateprep(np.array(input), wires=[0])], obs.diagonalizing_gates())
         res = dev.expval(obs)
 
         assert np.isclose(res, expected_output, atol=tol, rtol=0)
@@ -563,15 +563,16 @@ class TestVar:
             (qml.Identity, [1 / math.sqrt(2), -1 / math.sqrt(2)], 0),
         ],
     )
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
     def test_var_single_wire_no_parameters(
-        self, qubit_device, tol, operation, input, expected_output
+        self, qubit_device, tol, stateprep, operation, input, expected_output
     ):
         """Tests that variances are properly calculated for single-wire observables without parameters."""
         dev = qubit_device(wires=1)
         obs = operation(wires=[0])
 
         dev.reset()
-        dev.apply([qml.QubitStateVector(np.array(input), wires=[0])], obs.diagonalizing_gates())
+        dev.apply([stateprep(np.array(input), wires=[0])], obs.diagonalizing_gates())
         res = dev.var(obs)
 
         assert np.isclose(res, expected_output, atol=tol, rtol=0)
@@ -780,7 +781,10 @@ def circuit():
             ("CZ", [-1 / 2, -1 / 2]),
         ],
     )
-    def test_supported_gate_two_wires_no_parameters(self, qubit_device, tol, name, expected_output):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_supported_gate_two_wires_no_parameters(
+        self, qubit_device, tol, stateprep, name, expected_output
+    ):
         """Tests supported gates that act on two wires that are not parameterized"""
         dev = qubit_device(wires=2)
         op = getattr(qml.ops, name)
@@ -789,7 +793,7 @@ def test_supported_gate_two_wires_no_parameters(self, qubit_device, tol, name, e
 
         @qml.qnode(dev)
         def circuit():
-            qml.QubitStateVector(np.array([1 / 2, 0, 0, math.sqrt(3) / 2]), wires=[0, 1])
+            stateprep(np.array([1 / 2, 0, 0, math.sqrt(3) / 2]), wires=[0, 1])
             op(wires=[0, 1])
             return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
 
@@ -987,8 +991,9 @@ def circuit():
             ("ControlledPhaseShift", [math.pi], [-1 / 2, -1 / 2]),
         ],
     )
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
     def test_supported_gate_two_wires_with_parameters(
-        self, qubit_device, tol, name, par, expected_output
+        self, qubit_device, tol, stateprep, name, par, expected_output
     ):
         """Tests supported gates that act on two wires that are parameterized"""
         dev = qubit_device(wires=2)
@@ -998,7 +1003,7 @@ def test_supported_gate_two_wires_with_parameters(
 
         @qml.qnode(dev)
         def circuit():
-            qml.QubitStateVector(np.array([1 / 2, 0, 0, math.sqrt(3) / 2]), wires=[0, 1])
+            stateprep(np.array([1 / 2, 0, 0, math.sqrt(3) / 2]), wires=[0, 1])
             op(*par, wires=[0, 1])
             return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
 
@@ -1021,8 +1026,9 @@ def circuit():
             ("Hadamard", [1 / math.sqrt(2), 1 / math.sqrt(2)], 1 / math.sqrt(2)),
         ],
     )
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
     def test_supported_observable_single_wire_no_parameters(
-        self, qubit_device, tol, name, state, expected_output
+        self, qubit_device, tol, stateprep, name, state, expected_output
     ):
         """Tests supported observables on single wires without parameters."""
         dev = qubit_device(wires=1)
@@ -1032,7 +1038,7 @@ def test_supported_observable_single_wire_no_parameters(
 
         @qml.qnode(dev)
         def circuit():
-            qml.QubitStateVector(np.array(state), wires=[0])
+            stateprep(np.array(state), wires=[0])
             return qml.expval(obs(wires=[0]))
 
         assert np.isclose(circuit(), expected_output, atol=tol, rtol=0)
@@ -1045,8 +1051,9 @@ def circuit():
             ("Identity", [1 / math.sqrt(2), -1 / math.sqrt(2)], 1, []),
         ],
     )
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
     def test_supported_observable_single_wire_with_parameters(
-        self, qubit_device, tol, name, state, expected_output, par
+        self, qubit_device, tol, stateprep, name, state, expected_output, par
     ):
         """Tests supported observables on single wires with parameters."""
         dev = qubit_device(wires=1)
@@ -1056,7 +1063,7 @@ def test_supported_observable_single_wire_with_parameters(
 
         @qml.qnode(dev)
         def circuit():
-            qml.QubitStateVector(np.array(state), wires=[0])
+            stateprep(np.array(state), wires=[0])
             return qml.expval(obs(*par, wires=[0]))
 
         assert np.isclose(circuit(), expected_output, atol=tol, rtol=0)

tests/test_comparison.py

@@ -254,7 +254,10 @@ def circuit():
     @pytest.mark.parametrize("wires", range(1, 17))
     @pytest.mark.parametrize("num_threads", [1, 2])
     @pytest.mark.skipif(not ld._CPP_BINARY_AVAILABLE, reason="Lightning binary required")
-    def test_n_qubit_circuit(self, monkeypatch, wires, lightning_dev_version, num_threads):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_n_qubit_circuit(
+        self, monkeypatch, stateprep, wires, lightning_dev_version, num_threads
+    ):
         """Test an n-qubit circuit"""
         monkeypatch.setenv("OMP_NUM_THREADS", str(num_threads))
 
@@ -265,7 +268,7 @@ def test_n_qubit_circuit(self, monkeypatch, wires, lightning_dev_version, num_th
         def circuit():
             """Prepares the equal superposition state and then applies StronglyEntanglingLayers
             and concludes with a simple PauliZ measurement"""
-            qml.QubitStateVector(vec, wires=range(wires))
+            stateprep(vec, wires=range(wires))
             qml.StronglyEntanglingLayers(w, wires=range(wires))
             return qml.expval(qml.PauliZ(0))