Reintroduce QubitStateVector.

pennylane_lightning_gpu/lightning_gpu.py

@@ -141,6 +141,7 @@ def _mebibytesToBytes(mebibytes):
 allowed_operations = {
     "Identity",
     "BasisState",
+    "QubitStateVector",
     "StatePrep",
     "QubitUnitary",
     "ControlledQubitUnitary",

pennylane_lightning_gpu/src/algorithms/AdjointDiffGPU.hpp

@@ -562,7 +562,7 @@ template <class T = double> class AdjointJacobianGPU {
             PL_ABORT_IF(ops.getOpsParams()[op_idx].size() > 1,
                         "The operation is not supported using the adjoint "
                         "differentiation method");
-            if ((ops_name[op_idx] == "StatePrep") ||
+            if ((ops_name[op_idx] == "QubitStateVector") || (ops_name[op_idx] == "StatePrep") ||
                 (ops_name[op_idx] == "BasisState")) {
                 continue;
             }

pennylane_lightning_gpu/src/algorithms/AdjointDiffGPUMPI.hpp

@@ -302,7 +302,7 @@ class AdjointJacobianGPUMPI {
                 PL_ABORT_IF(ops.getOpsParams()[op_idx].size() > 1,
                             "The operation is not supported using the adjoint "
                             "differentiation method");
-                if ((ops_name[op_idx] == "StatePrep") ||
+                if ((ops_name[op_idx] == "QubitStateVector") || (ops_name[op_idx] == "StatePrep") ||
                     (ops_name[op_idx] == "BasisState")) {
                     continue;
                 }
@@ -398,7 +398,7 @@ class AdjointJacobianGPUMPI {
             PL_ABORT_IF(ops.getOpsParams()[op_idx].size() > 1,
                         "The operation is not supported using the adjoint "
                         "differentiation method");
-            if ((ops_name[op_idx] == "StatePrep") ||
+            if ((ops_name[op_idx] == "QubitStateVector") || (ops_name[op_idx] == "StatePrep") ||
                 (ops_name[op_idx] == "BasisState")) {
                 continue;
             }

tests/test_adjoint_jacobian.py

@@ -192,11 +192,12 @@ def test_unsupported_hermitian_expectation(self, dev_gpu):
 
     @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, tol, dev_cpu, dev_gpu):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_pauli_rotation_gradient(self, stateprep, G, theta, tol, dev_cpu, dev_gpu):
         """Tests that the automatic gradients of Pauli rotations are correct."""
 
         with qml.tape.QuantumTape() as tape:
-            qml.StatePrep(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
+            stateprep(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
             G(theta, wires=[0])
             qml.expval(qml.PauliZ(0))
 
@@ -208,13 +209,14 @@ def test_pauli_rotation_gradient(self, G, theta, tol, dev_cpu, dev_gpu):
         assert np.allclose(calculated_val, expected_val, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
-    def test_Rot_gradient(self, theta, tol, dev_cpu, dev_gpu):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_Rot_gradient(self, stateprep, theta, tol, dev_cpu, dev_gpu):
         """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.StatePrep(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
+            stateprep(np.array([1.0, -1.0]) / np.sqrt(2), wires=0)
             qml.Rot(*params, wires=[0])
             qml.expval(qml.PauliZ(0))
 

tests/test_apply.py

@@ -220,10 +220,10 @@ def test_apply_operation_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.StatePrep, [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]),
             (
@@ -477,17 +477,18 @@ def test_apply_operation_two_wires_with_parameters(
 
         assert np.allclose(state_vector, np.array(expected_output), atol=tol, rtol=0)
 
-    def test_apply_errors_qubit_state_vector(self, qubit_device_2_wires):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_apply_errors_qubit_state_vector(self, stateprep, qubit_device_2_wires):
         """Test that apply fails for incorrect state preparation, and > 2 qubit gates"""
         with pytest.raises(ValueError, match="Sum of amplitudes-squared does not equal one."):
-            qubit_device_2_wires.apply([qml.StatePrep(np.array([1, -1]), wires=[0])])
+            qubit_device_2_wires.apply([stateprep(np.array([1, -1]), wires=[0])])
 
         with pytest.raises(
             ValueError,
             match=r"State vector must have shape \(2\*\*wires,\) or \(batch_size, 2\*\*wires\).",
         ):
             p = np.array([1, 0, 1, 1, 0]) / np.sqrt(3)
-            qubit_device_2_wires.apply([qml.StatePrep(p, wires=[0, 1])])
+            qubit_device_2_wires.apply([stateprep(p, wires=[0, 1])])
 
         with pytest.raises(
             DeviceError,
@@ -497,7 +498,7 @@ def test_apply_errors_qubit_state_vector(self, qubit_device_2_wires):
             qubit_device_2_wires.apply(
                 [
                     qml.RZ(0.5, wires=[0]),
-                    qml.StatePrep(np.array([0, 1, 0, 0]), wires=[0, 1]),
+                    stateprep(np.array([0, 1, 0, 0]), wires=[0, 1]),
                 ]
             )
 
@@ -545,16 +546,17 @@ 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_1_wire, tol, operation, input, expected_output
+        self, stateprep, qubit_device_1_wire, tol, operation, input, expected_output
     ):
         """Tests that expectation values are properly calculated for single-wire observables without parameters."""
 
         obs = operation(wires=[0])
 
         qubit_device_1_wire.reset()
         qubit_device_1_wire.apply(
-            [qml.StatePrep(np.array(input), wires=[0])],
+            [stateprep(np.array(input), wires=[0])],
             rotations=obs.diagonalizing_gates(),
         )
         res = qubit_device_1_wire.expval(obs)
@@ -585,16 +587,17 @@ 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_1_wire, tol, operation, input, expected_output
+        self, stateprep, qubit_device_1_wire, tol, operation, input, expected_output
     ):
         """Tests that variances are properly calculated for single-wire observables without parameters."""
 
         obs = operation(wires=[0])
 
         qubit_device_1_wire.reset()
         qubit_device_1_wire.apply(
-            [qml.StatePrep(np.array(input), wires=[0])],
+            [stateprep(np.array(input), wires=[0])],
             rotations=obs.diagonalizing_gates(),
         )
         res = qubit_device_1_wire.var(obs)
@@ -749,8 +752,9 @@ def circuit():
             ("CZ", [-1 / 2, -1 / 2]),
         ],
     )
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
     def test_supported_gate_two_wires_no_parameters(
-        self, qubit_device_2_wires, tol, name, expected_output
+        self, stateprep, qubit_device_2_wires, tol, name, expected_output
     ):
         """Tests supported gates that act on two wires that are not parameterized"""
 
@@ -759,7 +763,7 @@ def test_supported_gate_two_wires_no_parameters(
 
         @qml.qnode(qubit_device_2_wires)
         def circuit():
-            qml.StatePrep(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))
 

tests/test_comparison.py

@@ -238,7 +238,8 @@ def circuit():
         assert np.allclose(default_state, lightninggpu_state)
 
     @pytest.mark.parametrize("wires", range(1, 17))
-    def test_n_qubit_circuit(self, wires, default_qubit_dev, lightning_gpu_qubit_dev):
+    @pytest.mark.parametrize("stateprep", [qml.QubitStateVector, qml.StatePrep])
+    def test_n_qubit_circuit(self, stateprep, wires, default_qubit_dev, lightning_gpu_qubit_dev):
         """Test an n-qubit circuit"""
 
         vec = np.array([1 + 0.0j] * (2**wires)) / np.sqrt(2**wires)
@@ -248,7 +249,7 @@ def test_n_qubit_circuit(self, wires, default_qubit_dev, lightning_gpu_qubit_dev
         def circuit():
             """Prepares the equal superposition state and then applies StronglyEntanglingLayers
             and concludes with a simple PauliZ measurement"""
-            qml.StatePrep(vec, wires=range(wires))
+            stateprep(vec, wires=range(wires))
             qml.StronglyEntanglingLayers(w, wires=range(wires))
             return qml.expval(qml.PauliZ(0))
 

tests/test_gates.py

@@ -96,7 +96,7 @@ def test_gate_unitary_correct(op, op_name):
     """Test if lightning.gpu correctly applies gates by reconstructing the unitary matrix and
     comparing to the expected version"""
 
-    if op_name in ("BasisState", "StatePrep"):
+    if op_name in ("BasisState", "QubitStateVector", "StatePrep"):
         pytest.skip("Skipping operation because it is a state preparation")
 
     if op == None:
@@ -131,7 +131,7 @@ def test_inverse_unitary_correct(op, op_name):
     """Test if lightning.gpu correctly applies inverse gates by reconstructing the unitary matrix
     and comparing to the expected version"""
 
-    if op_name in ("BasisState", "StatePrep"):
+    if op_name in ("BasisState", "QubitStateVector", "StatePrep"):
         pytest.skip("Skipping operation because it is a state preparation")
     if op == None:
         pytest.skip("Skipping operation.")