Repository: PennyLaneAI/pennylane

Issue:
Make sure all state preparation functions are compatible with all plugins
#### Issue description

We need to make sure all functions are supported in a compatible manner on all plugins. 

For example, qubit plugins other than `default.qubit` support the state prep function `BasisState` acting on ANY number of subsystems. `default.qubit` requires it to act on ALL wires. So the documentation for the function says `wires=Any` but produces an error when we use this function with the `default.qubit` plugin.

 For `QubitstateVector` function, it is probably the other way round. 

I will explicitly check all other state prep functions from gaussian plugin too. 



Discussion:
Thanks Aroosa! This is definitely something we want to double check and ensure consistency, especially where it concerns the default plugins.

Unfortunately, there may still be cases where different frameworks have different restrictions, for example Qiskit vs. PyQuil. In cases like this, we should probably pick the most permissive value of `num_wires`, and any plugin that is more restrictive should then raise its own error.
Operations we should double check:


- [ ] `BasisState`- ?
- [ ] `QubitStateVector` - `default.qubit` should be modified to allow state initialization on subsets of wires

- [ ] `QubitUnitary` - currently, `default.qubit` and Forest support arbitrary unitaries, Qiskit only suports single-qubit unitaries
> Unfortunately, there may still be cases where different frameworks have different restrictions, for example Qiskit vs. PyQuil. In cases like this, we should probably pick the most permissive value of `num_wires`, and any plugin that is more restrictive should then raise its own error.

We should explicitly mention this restriction somewhere. Maybe as a `Note` in the documentation. The plugin documentation only mention the added operations on top of pennylane operations so the pennylane documentation has to be complete. What do you think? 


My thinking is that the best place is the plugin readme. Every plugin should list the operations they support, as well as any changes/restrictions.

It will be too difficult to try and keep the PL documentation up to date with the rate the plugins are changing.
oh! that is a better suggestion :) 
Easier to update the readme if the plugin code ever changes as well, since they're in the same repo :)
@josh146 I have asked Angus to work on this issue as I have been a bit busy with the other project :) 

Resolving patch:
diff --git a/pennylane/ops/qubit.py b/pennylane/ops/qubit.py
--- a/pennylane/ops/qubit.py
+++ b/pennylane/ops/qubit.py
@@ -18,7 +18,7 @@
 # pylint:disable=abstract-method,arguments-differ
 import numpy as np
 
-from pennylane.operation import All, Any, Observable, Operation
+from pennylane.operation import Any, Observable, Operation
 from pennylane.templates.state_preparations import BasisStatePreparation, MottonenStatePreparation
 from pennylane.utils import OperationRecorder
 
@@ -701,7 +701,7 @@ class QubitUnitary(Operation):
 
     **Details:**
 
-    * Number of wires: The operation can act on any number of wires.
+    * Number of wires: Any (the operation can act on any number of wires)
     * Number of parameters: 1
     * Gradient recipe: None
 
@@ -726,7 +726,7 @@ class BasisState(Operation):
 
     **Details:**
 
-    * Number of wires: All (applied to the entire system)
+    * Number of wires: Any (the operation can act on any number of wires)
     * Number of parameters: 1
     * Gradient recipe: None (integer parameters not supported)
 
@@ -761,7 +761,7 @@ class QubitStateVector(Operation):
 
     **Details:**
 
-    * Number of wires: All (applied to the entire system)
+    * Number of wires: Any (the operation can act on any number of wires)
     * Number of parameters: 1
     * Gradient recipe: None
 
@@ -777,7 +777,7 @@ class QubitStateVector(Operation):
         wires (Sequence[int] or int): the wire(s) the operation acts on
     """
     num_params = 1
-    num_wires = All
+    num_wires = Any
     par_domain = "A"
     grad_method = None
 
diff --git a/pennylane/plugins/default_qubit.py b/pennylane/plugins/default_qubit.py
--- a/pennylane/plugins/default_qubit.py
+++ b/pennylane/plugins/default_qubit.py
@@ -27,7 +27,7 @@
 import numpy as np
 from scipy.linalg import eigh
 
-from pennylane import Device
+from pennylane import Device, DeviceError
 from pennylane.operation import Operation
 
 
@@ -317,37 +317,66 @@ def __init__(self, wires, *, shots=1000, analytic=True):
         self.analytic = analytic
 
         self._state = None
+        self._first_operation = True
 
     def pre_apply(self):
         self.reset()
 
     def apply(self, operation, wires, par):
+        # number of wires on device
+        n = self.num_wires
+
         if operation == 'QubitStateVector':
-            state = np.asarray(par[0], dtype=np.complex128)
-            if state.ndim == 1 and state.shape[0] == 2**self.num_wires:
-                self._state = state
+            input_state = np.asarray(par[0], dtype=np.complex128)
+
+            if not np.isclose(np.linalg.norm(input_state, 2), 1.0, atol=tolerance):
+                raise ValueError("Sum of amplitudes-squared does not equal one.")
+            n_state_vector = input_state.shape[0]
+            if not self._first_operation:
+                raise DeviceError("Operation {} cannot be used after other Operations have already been applied "
+                                  "on a {} device.".format(operation, self.short_name))
+            if input_state.ndim == 1 and n_state_vector == 2**len(wires):
+
+                # generate basis states on subset of qubits via the cartesian product
+                tuples = np.array(list(itertools.product([0, 1], repeat=len(wires))))
+
+                # get basis states to alter on full set of qubits
+                unravelled_nums = np.zeros((2 ** len(wires), n), dtype=int)
+                unravelled_nums[:, wires] = tuples
+
+                # get indices for which the state is changed to input state vector elements
+                nums = np.ravel_multi_index(unravelled_nums.T, [2] * n)
+                self._state = np.zeros_like(self._state)
+                self._state[nums] = input_state
             else:
-                raise ValueError('State vector must be of length 2**wires.')
-            if wires is not None and wires != [] and list(wires) != list(range(self.num_wires)):
-                raise ValueError("The default.qubit plugin can apply QubitStateVector only to all of the {} wires.".format(self.num_wires))
+                raise ValueError("State vector must be of length 2**wires.")
+            self._first_operation = False
             return
         if operation == 'BasisState':
-            n = len(par[0])
-            # get computational basis state number
-            if n > self.num_wires or not (set(par[0]) == {0, 1} or set(par[0]) == {0} or set(par[0]) == {1}):
-                raise ValueError("BasisState parameter must be an array of 0 or 1 integers of length at most {}.".format(self.num_wires))
-            if wires is not None and wires != [] and list(wires) != list(range(self.num_wires)):
-                raise ValueError("The default.qubit plugin can apply BasisState only to all of the {} wires.".format(self.num_wires))
+            # length of basis state parameter
+            n_basis_state = len(par[0])
+
+            if not set(par[0]).issubset({0, 1}):
+                raise ValueError("BasisState parameter must consist of 0 or 1 integers.")
+            if n_basis_state != len(wires):
+                raise ValueError("BasisState parameter and wires must be of equal length.")
+            if not self._first_operation:
+                raise DeviceError("Operation {} cannot be used after other Operations have already been applied "
+                                  "on a {} device.".format(operation, self.short_name))
 
-            num = int(np.sum(np.array(par[0])*2**np.arange(n-1, -1, -1)))
+
+            # get computational basis state number
+            num = int(np.dot(par[0], 2**(n - 1 - np.array(wires))))
 
             self._state = np.zeros_like(self._state)
             self._state[num] = 1.
+            self._first_operation = False
             return
 
         A = self._get_operator_matrix(operation, par)
 
         self._state = self.mat_vec_product(A, self._state, wires)
+        self._first_operation = False
 
     def mat_vec_product(self, mat, vec, wires):
         r"""Apply multiplication of a matrix to subsystems of the quantum state.
@@ -482,6 +511,7 @@ def reset(self):
         # init the state vector to |00..0>
         self._state = np.zeros(2**self.num_wires, dtype=complex)
         self._state[0] = 1
+        self._first_operation = True
 
     @property
     def operations(self):


Test patch:
diff --git a/tests/test_default_qubit.py b/tests/test_default_qubit.py
--- a/tests/test_default_qubit.py
+++ b/tests/test_default_qubit.py
@@ -20,7 +20,7 @@
 
 import pytest
 import pennylane as qml
-from pennylane import numpy as np
+from pennylane import numpy as np, DeviceError
 from pennylane.operation import Operation
 from pennylane.plugins.default_qubit import (CRot3, CRotx, CRoty, CRotz,
                                              Rot3, Rotx, Roty, Rotz,
@@ -452,20 +452,21 @@ def test_apply_operation_three_wires_no_parameters(self, qubit_device_3_wires, t
         assert np.allclose(qubit_device_3_wires._state, np.array(expected_output), atol=tol, rtol=0)
 
 
-    @pytest.mark.parametrize("name,input,expected_output,par", [
-        ("BasisState", [1, 0, 0, 0], [0, 0, 1, 0], [[1, 0]]),
-        ("BasisState", [1/math.sqrt(2), 0, 1/math.sqrt(2), 0], [0, 0, 1, 0], [[1, 0]]),
-        ("BasisState", [1/math.sqrt(2), 0, 1/math.sqrt(2), 0], [0, 0, 0, 1], [[1, 1]]),
-        ("QubitStateVector", [1, 0, 0, 0], [0, 0, 1, 0], [[0, 0, 1, 0]]),
-        ("QubitStateVector", [1/math.sqrt(2), 0, 1/math.sqrt(2), 0], [0, 0, 1, 0], [[0, 0, 1, 0]]),
-        ("QubitStateVector", [1/math.sqrt(2), 0, 1/math.sqrt(2), 0], [0, 0, 0, 1], [[0, 0, 0, 1]]),
-        ("QubitStateVector", [1, 0, 0, 0], [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)]]),
-        ("QubitStateVector", [1, 0, 0, 0], [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)]]),
+    @pytest.mark.parametrize("name,expected_output,par", [
+        ("BasisState", [0, 0, 1, 0], [[1, 0]]),
+        ("BasisState", [0, 0, 1, 0], [[1, 0]]),
+        ("BasisState", [0, 0, 0, 1], [[1, 1]]),
+        ("QubitStateVector", [0, 0, 1, 0], [[0, 0, 1, 0]]),
+        ("QubitStateVector", [0, 0, 1, 0], [[0, 0, 1, 0]]),
+        ("QubitStateVector", [0, 0, 0, 1], [[0, 0, 0, 1]]),
+        ("QubitStateVector", [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)]]),
+        ("QubitStateVector", [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)]]),
     ])
-    def test_apply_operation_state_preparation(self, qubit_device_2_wires, tol, name, input, expected_output, par):
-        """Tests that applying state preparation operations yield the expected output."""
+    def test_apply_operation_state_preparation(self, qubit_device_2_wires, tol, name, expected_output, par):
+        """Tests that applying an operation yields the expected output state for single wire
+           operations that have no parameters."""
 
-        qubit_device_2_wires._state = np.array(input)
+        qubit_device_2_wires.reset()
         qubit_device_2_wires.apply(name, wires=[0, 1], par=par)
 
         assert np.allclose(qubit_device_2_wires._state, np.array(expected_output), atol=tol, rtol=0)
@@ -571,27 +572,51 @@ def test_apply_operation_two_wires_with_parameters(self, qubit_device_2_wires, t
 
         assert np.allclose(qubit_device_2_wires._state, np.array(expected_output), atol=tol, rtol=0)
 
-    def test_apply_errors(self, qubit_device_2_wires):
+    def test_apply_errors_qubit_state_vector(self, 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("QubitStateVector", wires=[0], par=[np.array([1, -1])])
 
         with pytest.raises(
             ValueError,
             match=r"State vector must be of length 2\*\*wires."
         ):
-            p = [np.array([1, 0, 1, 1, 1]) / np.sqrt(3)]
+            p = np.array([1, 0, 1, 1, 0]) / np.sqrt(3)
             qubit_device_2_wires.apply("QubitStateVector", wires=[0, 1], par=[p])
 
+        with pytest.raises(
+            DeviceError,
+            match="Operation QubitStateVector cannot be used after other Operations have already been applied "
+                                  "on a default.qubit device."
+        ):
+            qubit_device_2_wires.reset()
+            qubit_device_2_wires.apply("RZ", wires=[0], par=[0.5])
+            qubit_device_2_wires.apply("QubitStateVector", wires=[0, 1], par=[np.array([0, 1, 0, 0])])
+
+    def test_apply_errors_basis_state(self, qubit_device_2_wires):
         with pytest.raises(
             ValueError,
-            match="BasisState parameter must be an array of 0 or 1 integers of length at most 2."
+            match="BasisState parameter must consist of 0 or 1 integers."
         ):
             qubit_device_2_wires.apply("BasisState", wires=[0, 1], par=[np.array([-0.2, 4.2])])
 
         with pytest.raises(
             ValueError,
-            match="The default.qubit plugin can apply BasisState only to all of the 2 wires."
+            match="BasisState parameter and wires must be of equal length."
+        ):
+            qubit_device_2_wires.apply("BasisState", wires=[0], par=[np.array([0, 1])])
+
+        with pytest.raises(
+            DeviceError,
+            match="Operation BasisState cannot be used after other Operations have already been applied "
+                                  "on a default.qubit device."
         ):
-            qubit_device_2_wires.apply("BasisState", wires=[0, 1, 2], par=[np.array([0, 1])])
+            qubit_device_2_wires.reset()
+            qubit_device_2_wires.apply("RZ", wires=[0], par=[0.5])
+            qubit_device_2_wires.apply("BasisState", wires=[0, 1], par=[[1, 1]])
 
 
 class TestExpval:
@@ -1058,6 +1083,71 @@ def circuit():
 
         assert np.allclose(circuit(), expected_output, atol=tol, rtol=0)
 
+    # This test is ran with two Z expvals
+    @pytest.mark.parametrize("name,par,wires,expected_output", [
+        ("BasisState", [1, 1], [0, 1], [-1, -1]),
+        ("BasisState", [1], [0], [-1, 1]),
+        ("BasisState", [1], [1], [1, -1])
+    ])
+    def test_basis_state_2_qubit_subset(self, qubit_device_2_wires, tol, name, par, wires, expected_output):
+        """Tests qubit basis state preparation on subsets of qubits"""
+
+        op = getattr(qml.ops, name)
+
+        @qml.qnode(qubit_device_2_wires)
+        def circuit():
+            op(np.array(par), wires=wires)
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
+
+        assert np.allclose(circuit(), expected_output, atol=tol, rtol=0)
+
+    # This test is run with two expvals
+    @pytest.mark.parametrize("name,par,wires,expected_output", [
+        ("QubitStateVector", [0, 1], [1], [1, -1]),
+        ("QubitStateVector", [0, 1], [0], [-1, 1]),
+        ("QubitStateVector", [1./np.sqrt(2), 1./np.sqrt(2)], [1], [1, 0]),
+        ("QubitStateVector", [1j/2., np.sqrt(3)/2.], [1], [1, -0.5]),
+        ("QubitStateVector", [(2-1j)/3., 2j/3.], [0], [1/9., 1])
+    ])
+    def test_state_vector_2_qubit_subset(self, qubit_device_2_wires, tol, name, par, wires, expected_output):
+        """Tests qubit state vector preparation on subsets of 2 qubits"""
+
+        op = getattr(qml.ops, name)
+
+        par = np.array(par)
+
+        @qml.qnode(qubit_device_2_wires)
+        def circuit():
+            op(par, wires=wires)
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
+
+        assert np.allclose(circuit(), expected_output, atol=tol, rtol=0)
+
+    # This test is run with three expvals
+    @pytest.mark.parametrize("name,par,wires,expected_output", [
+        ("QubitStateVector", [1j/np.sqrt(10), (1-2j)/np.sqrt(10), 0, 0, 0, 2/np.sqrt(10), 0, 0],
+         [0, 1, 2], [1/5., 1., -4/5.]),
+        ("QubitStateVector", [1/np.sqrt(2), 0, 0, 1/np.sqrt(2)], [0, 2], [0., 1., 0.]),
+        ("QubitStateVector", [1 / np.sqrt(2), 0, 0, 1 / np.sqrt(2)], [0, 1], [0., 0., 1.]),
+        ("QubitStateVector", [0, 1, 0, 0, 0, 0, 0, 0], [2, 1, 0], [-1., 1., 1.]),
+        ("QubitStateVector", [0, 1j, 0, 0, 0, 0, 0, 0], [0, 2, 1], [1., -1., 1.]),
+        ("QubitStateVector", [0, 1/np.sqrt(2), 0, 1/np.sqrt(2)], [1, 0], [-1., 0., 1.]),
+        ("QubitStateVector", [0, 1 / np.sqrt(2), 0, 1 / np.sqrt(2)], [0, 1], [0., -1., 1.])
+    ])
+    def test_state_vector_3_qubit_subset(self, qubit_device_3_wires, tol, name, par, wires, expected_output):
+        """Tests qubit state vector preparation on subsets of 3 qubits"""
+
+        op = getattr(qml.ops, name)
+
+        par = np.array(par)
+
+        @qml.qnode(qubit_device_3_wires)
+        def circuit():
+            op(par, wires=wires)
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)), qml.expval(qml.PauliZ(2))
+
+        assert np.allclose(circuit(), expected_output, atol=tol, rtol=0)
+
     # This test is ran on the state |0> with one Z expvals
     @pytest.mark.parametrize("name,par,expected_output", [
         ("PhaseShift", [math.pi/2], 1),