[TEMPLATES] Rewrite state preparation templates as operations (#1190)

* rewrite arbitrary state prep

* rewrite arbitrary state prep

* backup

* all tests passing

* improve docstrings

* polish

* black

* fix qasm

* black

* adjust expansion depths

* black

* Update pennylane/templates/state_preparations/arbitrary_state_preparation.py

Co-authored-by: Christina Lee <christina@xanadu.ai>

* applied code review

* fix test imports in embeddings tests

* improve test

* updated changelog

Co-authored-by: Christina Lee <christina@xanadu.ai>

.github/CHANGELOG.md

@@ -537,14 +537,16 @@
   1: ──RY(1.35)──╰X──RY(0.422)──╰X──┤   
   ```
 
-- The embedding templates, as well as `BasicEntanglerLayers`, are now classes inheriting 
+- The embedding and state preparation templates, 
+  as well as `BasicEntanglerLayers`, are now classes inheriting 
   from `Operation`, and define the ansatz in their `expand()` method. This 
   change does not affect the user interface.
   [(#1138)](https://github.com/PennyLaneAI/pennylane/pull/1138)
   [(#1156)](https://github.com/PennyLaneAI/pennylane/pull/1156)
 
-  For convenience, `BasicEntanglerLayers` has a method that returns the shape of the 
-  trainable parameter tensor, i.e.,
+  For convenience, some templates now have a method that returns the expected 
+  shape of the trainable parameter tensor, which can be used to create 
+  random tensors.
 
   ```python
   shape = qml.templates.BasicEntanglerLayers.shape(n_layers=2, n_wires=4)

pennylane/tape/tape.py

@@ -1103,7 +1103,7 @@ def to_openqasm(self, wires=None, rotations=True):
                     op.inv()
 
         # decompose the queue
-        operations = tape.expand(stop_at=lambda obj: obj.name in OPENQASM_GATES).operations
+        operations = tape.expand(depth=2, stop_at=lambda obj: obj.name in OPENQASM_GATES).operations
 
         # create the QASM code representing the operations
         for op in operations:

pennylane/templates/state_preparations/arbitrary_state_preparation.py

@@ -12,28 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 r"""
-Contains the ``ArbitraryStatePreparation`` template.
+Contains the ArbitraryStatePreparation template.
 """
+# pylint: disable=trailing-comma-tuple
 import functools
 import pennylane as qml
-from pennylane.templates.decorator import template
-from pennylane.wires import Wires
-
-
-def _preprocess(weights, wires):
-    """Validate and pre-process inputs as follows:
-
-    * Check the shape of the weights tensor.
-
-    Args:
-        weights (tensor_like): trainable parameters of the template
-        wires (Wires): wires that template acts on
-    """
-    shape = qml.math.shape(weights)
-    if shape != (2 ** (len(wires) + 1) - 2,):
-        raise ValueError(
-            f"Weights tensor must be of shape {(2 ** (len(wires) + 1) - 2,)}; got {shape}."
-        )
+from pennylane.operation import Operation, AnyWires
 
 
 @functools.lru_cache()
@@ -60,8 +44,7 @@ def _state_preparation_pauli_words(num_wires):
     return single_qubit_words + multi_qubit_words
 
 
-@template
-def ArbitraryStatePreparation(weights, wires):
+class ArbitraryStatePreparation(Operation):
     """Implements an arbitrary state preparation on the specified wires.
 
     An arbitrary state on :math:`n` wires is parametrized by :math:`2^{n+1} - 2`
@@ -79,19 +62,53 @@ def ArbitraryStatePreparation(weights, wires):
 
         @qml.qnode(dev)
         def vqe(weights):
-            qml.ArbitraryStatePreparations(weights, wires=[0, 1, 2, 3])
+            qml.templates.ArbitraryStatePreparation(weights, wires=[0, 1, 2, 3])
 
             return qml.expval(qml.Hermitian(H, wires=[0, 1, 2, 3]))
 
+    The shape of the weights parameter can be computed as follows:
+
+    .. code-block:: python
+
+        shape = qml.templates.ArbitraryStatePreparation.shape(n_wires=4)
+
+
     Args:
-        weights (tensor_like): The angles of the Pauli word rotations, needs to have length :math:`2^(n+1) - 2`
+        weights (tensor_like): Angles of the Pauli word rotations. Needs to have length :math:`2^(n+1) - 2`
             where :math:`n` is the number of wires the template acts upon.
-        wires (Iterable or Wires): Wires that the template acts on. Accepts an iterable of numbers or strings, or
-            a Wires object.
+        wires (Iterable): wires that the template acts on
     """
 
-    wires = Wires(wires)
-    _preprocess(weights, wires)
+    num_params = 1
+    num_wires = AnyWires
+    par_domain = "A"
+
+    def __init__(self, weights, wires, do_queue=True):
+
+        shape = qml.math.shape(weights)
+        if shape != (2 ** (len(wires) + 1) - 2,):
+            raise ValueError(
+                f"Weights tensor must be of shape {(2 ** (len(wires) + 1) - 2,)}; got {shape}."
+            )
+
+        super().__init__(weights, wires=wires, do_queue=do_queue)
+
+    def expand(self):
+
+        with qml.tape.QuantumTape() as tape:
+            for i, pauli_word in enumerate(_state_preparation_pauli_words(len(self.wires))):
+                qml.PauliRot(self.parameters[0][i], pauli_word, wires=self.wires)
+
+        return tape
+
+    @staticmethod
+    def shape(n_wires):
+        r"""Returns the required shape for the weight tensor.
+
+        Args:
+                n_wires (int): number of wires
 
-    for i, pauli_word in enumerate(_state_preparation_pauli_words(len(wires))):
-        qml.PauliRot(weights[i], pauli_word, wires=wires)
+        Returns:
+            tuple[int]: shape
+        """
+        return (2 ** (n_wires + 1) - 2,)

pennylane/templates/state_preparations/basis.py

@@ -12,71 +12,69 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 r"""
-Contains the ``BasisStatePreparation`` template.
+Contains the BasisStatePreparation template.
 """
 
 import pennylane as qml
-from pennylane.templates.decorator import template
-from pennylane.wires import Wires
+from pennylane.operation import Operation, AnyWires
 
 
-def _preprocess(basis_state, wires):
-    """Validate and pre-process inputs as follows:
+class BasisStatePreparation(Operation):
+    r"""
+    Prepares a basis state on the given wires using a sequence of Pauli-X gates.
 
-    * Check the shape of the basis state.
-    * Cast basis state to a numpy array.
+    .. warning::
 
-    Args:
-        basis_state (tensor_like): basis state to prepare
-        wires (Wires): wires that template acts on
+        ``basis_state`` influences the circuit architecture and is therefore incompatible with
+        gradient computations.
 
-    Returns:
-        array: preprocessed basis state
-    """
-    shape = qml.math.shape(basis_state)
+    **Example**
 
-    if len(shape) != 1:
-        raise ValueError(f"Basis state must be one-dimensional; got shape {shape}.")
+    .. code-block:: python
 
-    n_bits = shape[0]
-    if n_bits != len(wires):
-        raise ValueError(f"Basis state must be of length {len(wires)}; got length {n_bits}.")
+        dev = qml.device("default.qubit", wires=4)
 
-    basis_state = list(qml.math.toarray(basis_state))
+        @qml.qnode(dev)
+        def circuit(basis_state):
+            qml.templates.BasisStatePreparation(basis_state, wires=range(4))
 
-    if not all(bit in [0, 1] for bit in basis_state):
-        raise ValueError(f"Basis state must only consist of 0s and 1s; got {basis_state}")
+            return [qml.expval(qml.PauliZ(wires=i)) for i in range(4)]
 
-    # we return the input as a list of values, since
-    # it is not differentiable
-    return basis_state
+        basis_state = [0, 1, 1, 0]
+        print(circuit(basis_state)) # [ 1. -1. -1.  1.]
 
+    Args:
+        basis_state (array): Input array of shape ``(n,)``, where n is the number of wires
+            the state preparation acts on.
+        wires (Iterable): wires that the template acts on
+    """
+    num_params = 1
+    num_wires = AnyWires
+    par_domain = "A"
 
-@template
-def BasisStatePreparation(basis_state, wires):
-    r"""
-    Prepares a basis state on the given wires using a sequence of Pauli X gates.
+    def __init__(self, basis_state, wires, do_queue=True):
 
-    .. warning::
+        shape = qml.math.shape(basis_state)
 
-        ``basis_state`` influences the circuit architecture and is therefore incompatible with
-        gradient computations. Ensure that ``basis_state`` is not passed to the qnode by positional
-        arguments.
+        if len(shape) != 1:
+            raise ValueError(f"Basis state must be one-dimensional; got shape {shape}.")
 
-    Args:
-        basis_state (array): Input array of shape ``(N,)``, where N is the number of wires
-            the state preparation acts on. ``N`` must be smaller or equal to the total
-            number of wires of the device.
-        wires (Iterable or Wires): Wires that the template acts on. Accepts an iterable of numbers or strings, or
-            a Wires object.
-
-    Raises:
-        ValueError: if inputs do not have the correct format
-    """
-    wires = Wires(wires)
+        n_bits = shape[0]
+        if n_bits != len(wires):
+            raise ValueError(f"Basis state must be of length {len(wires)}; got length {n_bits}.")
+
+        # we can extract a list here, because embedding is not differentiable
+        basis_state = list(qml.math.toarray(basis_state))
+
+        if not all(bit in [0, 1] for bit in basis_state):
+            raise ValueError(f"Basis state must only consist of 0s and 1s; got {basis_state}")
+
+        super().__init__(basis_state, wires=wires, do_queue=do_queue)
 
-    basis_state = _preprocess(basis_state, wires)
+    def expand(self):
 
-    for wire, state in zip(wires, basis_state):
-        if state == 1:
-            qml.PauliX(wire)
+        with qml.tape.QuantumTape() as tape:
+            for wire, state in zip(self.wires, self.parameters[0]):
+                if state == 1:
+                    qml.PauliX(wire)
+        return tape