We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Unexpected behavior with merge_amplitude_embedding and batched input
### Expected behavior

When using merge_amplitude_embedding to combine two separate embeddings, the output should be the same as two separate MottonenStatePreparations.

### Actual behavior

Instead, what I observe is that the output not only doesn't match the batch size, but it's also incorrect.


### Additional information

Removing `@qml.transforms.merge_amplitude_embedding` solves the issue, but I'm not sure it's intended.

### Source code

```shell
import pennylane as qml
import torch

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev, interface="torch")
def circuit1(a, b): 
    qml.MottonenStatePreparation(a, wires=0)
    qml.MottonenStatePreparation(b, wires=1)
    return qml.probs(wires=range(2))

a = torch.randn((3,2))
b = torch.randn((3,2))
a = a/a.norm(dim=1, keepdim=True)
b = b/b.norm(dim=1, keepdim=True)

print('a:')
print(a)
print('b:')
print(b)

print('Batched call to circuit1')
print(circuit1(a, b).detach().numpy())

# Output:
# a:
# tensor([[ 0.4929,  0.8701],
#         [ 0.7628,  0.6466],
#         [-0.6488, -0.7610]])
# b:
# tensor([[-0.6827, -0.7307],
#         [ 0.6346, -0.7728],
#         [ 0.2947, -0.9556]])
# Batched call to circuit1
# [[0.11325859 0.12973125 0.35284562 0.40416454]
#  [0.23433313 0.34756744 0.16836992 0.24972952]
#  [0.03655629 0.38435813 0.0502934  0.52879218]]


@qml.transforms.merge_amplitude_embedding
@qml.qnode(dev, interface="torch")
def circuit2(a, b): 
    qml.AmplitudeEmbedding(a, wires=0)
    qml.AmplitudeEmbedding(b, wires=1)
    return qml.probs(wires=range(2))

print('Batched call to circuit2')
print(circuit2(a, b).detach().numpy())

print('Repeated call to circuit2')
print(circuit2(a[0], b[0]).detach().numpy())
print(circuit2(a[1], b[1]).detach().numpy())
print(circuit2(a[2], b[2]).detach().numpy())


# Output:
# Batched call to circuit2
# [1. 0. 0. 0.]
# Repeated call to circuit2
# [1. 0. 0. 0.]
# [1. 0. 0. 0.]
# [1. 0. 0. 0.]
```


### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.33.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /home/tiblias/miniconda3/envs/qc-gpu/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning, PennyLane-Lightning-GPU, PennyLane-qiskit, pennylane-qulacs

Platform info:           Linux-6.2.0-36-generic-x86_64-with-glibc2.35
Python version:          3.10.12
Numpy version:           1.23.5
Scipy version:           1.11.2
Installed devices:
- default.gaussian (PennyLane-0.33.0)
- default.mixed (PennyLane-0.33.0)
- default.qubit (PennyLane-0.33.0)
- default.qubit.autograd (PennyLane-0.33.0)
- default.qubit.jax (PennyLane-0.33.0)
- default.qubit.legacy (PennyLane-0.33.0)
- default.qubit.tf (PennyLane-0.33.0)
- default.qubit.torch (PennyLane-0.33.0)
- default.qutrit (PennyLane-0.33.0)
- null.qubit (PennyLane-0.33.0)
- lightning.qubit (PennyLane-Lightning-0.33.1)
- lightning.gpu (PennyLane-Lightning-GPU-0.31.0)
- qiskit.aer (PennyLane-qiskit-0.31.0)
- qiskit.basicaer (PennyLane-qiskit-0.31.0)
- qiskit.ibmq (PennyLane-qiskit-0.31.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.31.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.31.0)
- qiskit.remote (PennyLane-qiskit-0.31.0)
- qulacs.simulator (pennylane-qulacs-0.29.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pennylane/transforms/optimization/merge_amplitude_embedding.py`
Content:
```
1 # Copyright 2018-2021 Xanadu Quantum Technologies Inc.
2 
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """Transform for merging AmplitudeEmbedding gates in a quantum circuit."""
15 from typing import Sequence, Callable
16 
17 from pennylane.transforms import transform
18 from pennylane.tape import QuantumTape
19 from pennylane import AmplitudeEmbedding
20 from pennylane._device import DeviceError
21 from pennylane.math import flatten, reshape
22 
23 
24 @transform
25 def merge_amplitude_embedding(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):
26     r"""Quantum function transform to combine amplitude embedding templates that act on different qubits.
27 
28     Args:
29         tape (QNode or QuantumTape or Callable): A quantum circuit.
30 
31     Returns:
32         qnode (QNode) or quantum function (Callable) or tuple[List[.QuantumTape], function]: The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
33 
34 
35     **Example**
36 
37     >>> dev = qml.device('default.qubit', wires=4)
38 
39     You can apply the transform directly on :class:`QNode`:
40 
41     .. code-block:: python
42 
43         @qml.transforms.merge_amplitude_embedding
44         @qml.qnode(device=dev)
45         def circuit():
46             qml.CNOT(wires = [0,1])
47             qml.AmplitudeEmbedding([0,1], wires = 2)
48             qml.AmplitudeEmbedding([0,1], wires = 3)
49             return qml.state()
50 
51     >>> circuit()
52     [1.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j]
53 
54     .. details::
55         :title: Usage Details
56 
57         You can also apply it on quantum function.
58 
59         .. code-block:: python
60 
61             def qfunc():
62                 qml.CNOT(wires = [0,1])
63                 qml.AmplitudeEmbedding([0,1], wires = 2)
64                 qml.AmplitudeEmbedding([0,1], wires = 3)
65                 return qml.state()
66 
67         The circuit before compilation will not work because of using two amplitude embedding.
68 
69         Using the transformation we can join the different amplitude embedding into a single one:
70 
71         >>> optimized_qfunc = qml.transforms.merge_amplitude_embedding(qfunc)
72         >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
73         >>> print(qml.draw(optimized_qnode)())
74         0: ─╭●──────────────────────┤  State
75         1: ─╰X──────────────────────┤  State
76         2: ─╭AmplitudeEmbedding(M0)─┤  State
77         3: ─╰AmplitudeEmbedding(M0)─┤  State
78         M0 =
79         [0.+0.j 0.+0.j 0.+0.j 1.+0.j]
80 
81     """
82     # Make a working copy of the list to traverse
83     list_copy = tape.operations.copy()
84     not_amplitude_embedding = []
85     visited_wires = set()
86     input_wires, input_vectors, input_batch_size = [], [], []
87     while len(list_copy) > 0:
88         current_gate = list_copy[0]
89         wires_set = set(current_gate.wires)
90 
91         # Check if the current gate is an AmplitudeEmbedding.
92         if not isinstance(current_gate, AmplitudeEmbedding):
93             not_amplitude_embedding.append(current_gate)
94             list_copy.pop(0)
95             visited_wires = visited_wires.union(wires_set)
96             continue
97 
98         # Check the qubits have not been used.
99         if len(visited_wires.intersection(wires_set)) > 0:
100             raise DeviceError(
101                 f"Operation {current_gate.name} cannot be used after other Operation applied in the same qubit "
102             )
103         input_wires.append(current_gate.wires)
104         input_vectors.append(current_gate.parameters[0])
105         input_batch_size.append(current_gate.batch_size)
106         list_copy.pop(0)
107         visited_wires = visited_wires.union(wires_set)
108 
109     if len(input_wires) > 0:
110         final_wires = input_wires[0]
111         final_vector = input_vectors[0]
112         final_batch_size = input_batch_size[0]
113 
114         # Merge all parameters and qubits into a single one.
115         for w, v, b in zip(input_wires[1:], input_vectors[1:], input_batch_size[1:]):
116             final_vector = final_vector[..., :, None] * v[..., None, :]
117             final_batch_size = final_batch_size or b
118             final_wires = final_wires + w
119 
120             if final_batch_size:
121                 final_vector = reshape(final_vector, (final_batch_size, -1))
122             else:
123                 final_vector = flatten(final_vector)
124 
125         AmplitudeEmbedding(final_vector, wires=final_wires)
126 
127     new_operations = []
128     for gate in not_amplitude_embedding:
129         new_operations.append(gate)
130 
131     new_tape = type(tape)(new_operations, tape.measurements, shots=tape.shots)
132 
133     def null_postprocessing(results):
134         """A postprocesing function returned by a transform that only converts the batch of results
135         into a result for a single ``QuantumTape``.
136         """
137         return results[0]
138 
139     return [new_tape], null_postprocessing
140
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pennylane/transforms/optimization/merge_amplitude_embedding.py b/pennylane/transforms/optimization/merge_amplitude_embedding.py
--- a/pennylane/transforms/optimization/merge_amplitude_embedding.py
+++ b/pennylane/transforms/optimization/merge_amplitude_embedding.py
@@ -19,6 +19,7 @@
 from pennylane import AmplitudeEmbedding
 from pennylane._device import DeviceError
 from pennylane.math import flatten, reshape
+from pennylane.queuing import QueuingManager
 
 
 @transform
@@ -79,19 +80,15 @@
         [0.+0.j 0.+0.j 0.+0.j 1.+0.j]
 
     """
-    # Make a working copy of the list to traverse
-    list_copy = tape.operations.copy()
-    not_amplitude_embedding = []
+    new_operations = []
     visited_wires = set()
     input_wires, input_vectors, input_batch_size = [], [], []
-    while len(list_copy) > 0:
-        current_gate = list_copy[0]
+    for current_gate in tape.operations:
         wires_set = set(current_gate.wires)
 
         # Check if the current gate is an AmplitudeEmbedding.
         if not isinstance(current_gate, AmplitudeEmbedding):
-            not_amplitude_embedding.append(current_gate)
-            list_copy.pop(0)
+            new_operations.append(current_gate)
             visited_wires = visited_wires.union(wires_set)
             continue
 
@@ -103,7 +100,6 @@
         input_wires.append(current_gate.wires)
         input_vectors.append(current_gate.parameters[0])
         input_batch_size.append(current_gate.batch_size)
-        list_copy.pop(0)
         visited_wires = visited_wires.union(wires_set)
 
     if len(input_wires) > 0:
@@ -122,11 +118,8 @@
             else:
                 final_vector = flatten(final_vector)
 
-        AmplitudeEmbedding(final_vector, wires=final_wires)
-
-    new_operations = []
-    for gate in not_amplitude_embedding:
-        new_operations.append(gate)
+        with QueuingManager.stop_recording():
+            new_operations.insert(0, AmplitudeEmbedding(final_vector, wires=final_wires))
 
     new_tape = type(tape)(new_operations, tape.measurements, shots=tape.shots)

{"golden_diff": "diff --git a/pennylane/transforms/optimization/merge_amplitude_embedding.py b/pennylane/transforms/optimization/merge_amplitude_embedding.py\n--- a/pennylane/transforms/optimization/merge_amplitude_embedding.py\n+++ b/pennylane/transforms/optimization/merge_amplitude_embedding.py\n@@ -19,6 +19,7 @@\n from pennylane import AmplitudeEmbedding\n from pennylane._device import DeviceError\n from pennylane.math import flatten, reshape\n+from pennylane.queuing import QueuingManager\n \n \n @transform\n@@ -79,19 +80,15 @@\n         [0.+0.j 0.+0.j 0.+0.j 1.+0.j]\n \n     \"\"\"\n-    # Make a working copy of the list to traverse\n-    list_copy = tape.operations.copy()\n-    not_amplitude_embedding = []\n+    new_operations = []\n     visited_wires = set()\n     input_wires, input_vectors, input_batch_size = [], [], []\n-    while len(list_copy) > 0:\n-        current_gate = list_copy[0]\n+    for current_gate in tape.operations:\n         wires_set = set(current_gate.wires)\n \n         # Check if the current gate is an AmplitudeEmbedding.\n         if not isinstance(current_gate, AmplitudeEmbedding):\n-            not_amplitude_embedding.append(current_gate)\n-            list_copy.pop(0)\n+            new_operations.append(current_gate)\n             visited_wires = visited_wires.union(wires_set)\n             continue\n \n@@ -103,7 +100,6 @@\n         input_wires.append(current_gate.wires)\n         input_vectors.append(current_gate.parameters[0])\n         input_batch_size.append(current_gate.batch_size)\n-        list_copy.pop(0)\n         visited_wires = visited_wires.union(wires_set)\n \n     if len(input_wires) > 0:\n@@ -122,11 +118,8 @@\n             else:\n                 final_vector = flatten(final_vector)\n \n-        AmplitudeEmbedding(final_vector, wires=final_wires)\n-\n-    new_operations = []\n-    for gate in not_amplitude_embedding:\n-        new_operations.append(gate)\n+        with QueuingManager.stop_recording():\n+            new_operations.insert(0, AmplitudeEmbedding(final_vector, wires=final_wires))\n \n     new_tape = type(tape)(new_operations, tape.measurements, shots=tape.shots)\n", "issue": "[BUG] Unexpected behavior with merge_amplitude_embedding and batched input\n### Expected behavior\n\nWhen using merge_amplitude_embedding to combine two separate embeddings, the output should be the same as two separate MottonenStatePreparations.\n\n### Actual behavior\n\nInstead, what I observe is that the output not only doesn't match the batch size, but it's also incorrect.\r\n\n\n### Additional information\n\nRemoving `@qml.transforms.merge_amplitude_embedding` solves the issue, but I'm not sure it's intended.\n\n### Source code\n\n```shell\nimport pennylane as qml\r\nimport torch\r\n\r\ndev = qml.device(\"default.qubit\", wires=2)\r\n\r\n@qml.qnode(dev, interface=\"torch\")\r\ndef circuit1(a, b): \r\n    qml.MottonenStatePreparation(a, wires=0)\r\n    qml.MottonenStatePreparation(b, wires=1)\r\n    return qml.probs(wires=range(2))\r\n\r\na = torch.randn((3,2))\r\nb = torch.randn((3,2))\r\na = a/a.norm(dim=1, keepdim=True)\r\nb = b/b.norm(dim=1, keepdim=True)\r\n\r\nprint('a:')\r\nprint(a)\r\nprint('b:')\r\nprint(b)\r\n\r\nprint('Batched call to circuit1')\r\nprint(circuit1(a, b).detach().numpy())\r\n\r\n# Output:\r\n# a:\r\n# tensor([[ 0.4929,  0.8701],\r\n#         [ 0.7628,  0.6466],\r\n#         [-0.6488, -0.7610]])\r\n# b:\r\n# tensor([[-0.6827, -0.7307],\r\n#         [ 0.6346, -0.7728],\r\n#         [ 0.2947, -0.9556]])\r\n# Batched call to circuit1\r\n# [[0.11325859 0.12973125 0.35284562 0.40416454]\r\n#  [0.23433313 0.34756744 0.16836992 0.24972952]\r\n#  [0.03655629 0.38435813 0.0502934  0.52879218]]\r\n\r\n\r\n@qml.transforms.merge_amplitude_embedding\r\n@qml.qnode(dev, interface=\"torch\")\r\ndef circuit2(a, b): \r\n    qml.AmplitudeEmbedding(a, wires=0)\r\n    qml.AmplitudeEmbedding(b, wires=1)\r\n    return qml.probs(wires=range(2))\r\n\r\nprint('Batched call to circuit2')\r\nprint(circuit2(a, b).detach().numpy())\r\n\r\nprint('Repeated call to circuit2')\r\nprint(circuit2(a[0], b[0]).detach().numpy())\r\nprint(circuit2(a[1], b[1]).detach().numpy())\r\nprint(circuit2(a[2], b[2]).detach().numpy())\r\n\r\n\r\n# Output:\r\n# Batched call to circuit2\r\n# [1. 0. 0. 0.]\r\n# Repeated call to circuit2\r\n# [1. 0. 0. 0.]\r\n# [1. 0. 0. 0.]\r\n# [1. 0. 0. 0.]\n```\n\n\n### Tracebacks\n\n_No response_\n\n### System information\n\n```shell\nName: PennyLane\r\nVersion: 0.33.0\r\nSummary: PennyLane is a Python quantum machine learning library by Xanadu Inc.\r\nHome-page: https://github.com/PennyLaneAI/pennylane\r\nAuthor: \r\nAuthor-email: \r\nLicense: Apache License 2.0\r\nLocation: /home/tiblias/miniconda3/envs/qc-gpu/lib/python3.10/site-packages\r\nRequires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions\r\nRequired-by: PennyLane-Lightning, PennyLane-Lightning-GPU, PennyLane-qiskit, pennylane-qulacs\r\n\r\nPlatform info:           Linux-6.2.0-36-generic-x86_64-with-glibc2.35\r\nPython version:          3.10.12\r\nNumpy version:           1.23.5\r\nScipy version:           1.11.2\r\nInstalled devices:\r\n- default.gaussian (PennyLane-0.33.0)\r\n- default.mixed (PennyLane-0.33.0)\r\n- default.qubit (PennyLane-0.33.0)\r\n- default.qubit.autograd (PennyLane-0.33.0)\r\n- default.qubit.jax (PennyLane-0.33.0)\r\n- default.qubit.legacy (PennyLane-0.33.0)\r\n- default.qubit.tf (PennyLane-0.33.0)\r\n- default.qubit.torch (PennyLane-0.33.0)\r\n- default.qutrit (PennyLane-0.33.0)\r\n- null.qubit (PennyLane-0.33.0)\r\n- lightning.qubit (PennyLane-Lightning-0.33.1)\r\n- lightning.gpu (PennyLane-Lightning-GPU-0.31.0)\r\n- qiskit.aer (PennyLane-qiskit-0.31.0)\r\n- qiskit.basicaer (PennyLane-qiskit-0.31.0)\r\n- qiskit.ibmq (PennyLane-qiskit-0.31.0)\r\n- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.31.0)\r\n- qiskit.ibmq.sampler (PennyLane-qiskit-0.31.0)\r\n- qiskit.remote (PennyLane-qiskit-0.31.0)\r\n- qulacs.simulator (pennylane-qulacs-0.29.0)\n```\n\n\n### Existing GitHub issues\n\n- [X] I have searched existing GitHub issues to make sure the issue does not already exist.\n", "before_files": [{"content": "# Copyright 2018-2021 Xanadu Quantum Technologies Inc.\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Transform for merging AmplitudeEmbedding gates in a quantum circuit.\"\"\"\nfrom typing import Sequence, Callable\n\nfrom pennylane.transforms import transform\nfrom pennylane.tape import QuantumTape\nfrom pennylane import AmplitudeEmbedding\nfrom pennylane._device import DeviceError\nfrom pennylane.math import flatten, reshape\n\n\n@transform\ndef merge_amplitude_embedding(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):\n    r\"\"\"Quantum function transform to combine amplitude embedding templates that act on different qubits.\n\n    Args:\n        tape (QNode or QuantumTape or Callable): A quantum circuit.\n\n    Returns:\n        qnode (QNode) or quantum function (Callable) or tuple[List[.QuantumTape], function]: The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.\n\n\n    **Example**\n\n    >>> dev = qml.device('default.qubit', wires=4)\n\n    You can apply the transform directly on :class:`QNode`:\n\n    .. code-block:: python\n\n        @qml.transforms.merge_amplitude_embedding\n        @qml.qnode(device=dev)\n        def circuit():\n            qml.CNOT(wires = [0,1])\n            qml.AmplitudeEmbedding([0,1], wires = 2)\n            qml.AmplitudeEmbedding([0,1], wires = 3)\n            return qml.state()\n\n    >>> circuit()\n    [1.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j]\n\n    .. details::\n        :title: Usage Details\n\n        You can also apply it on quantum function.\n\n        .. code-block:: python\n\n            def qfunc():\n                qml.CNOT(wires = [0,1])\n                qml.AmplitudeEmbedding([0,1], wires = 2)\n                qml.AmplitudeEmbedding([0,1], wires = 3)\n                return qml.state()\n\n        The circuit before compilation will not work because of using two amplitude embedding.\n\n        Using the transformation we can join the different amplitude embedding into a single one:\n\n        >>> optimized_qfunc = qml.transforms.merge_amplitude_embedding(qfunc)\n        >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)\n        >>> print(qml.draw(optimized_qnode)())\n        0: \u2500\u256d\u25cf\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2524  State\n        1: \u2500\u2570X\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2524  State\n        2: \u2500\u256dAmplitudeEmbedding(M0)\u2500\u2524  State\n        3: \u2500\u2570AmplitudeEmbedding(M0)\u2500\u2524  State\n        M0 =\n        [0.+0.j 0.+0.j 0.+0.j 1.+0.j]\n\n    \"\"\"\n    # Make a working copy of the list to traverse\n    list_copy = tape.operations.copy()\n    not_amplitude_embedding = []\n    visited_wires = set()\n    input_wires, input_vectors, input_batch_size = [], [], []\n    while len(list_copy) > 0:\n        current_gate = list_copy[0]\n        wires_set = set(current_gate.wires)\n\n        # Check if the current gate is an AmplitudeEmbedding.\n        if not isinstance(current_gate, AmplitudeEmbedding):\n            not_amplitude_embedding.append(current_gate)\n            list_copy.pop(0)\n            visited_wires = visited_wires.union(wires_set)\n            continue\n\n        # Check the qubits have not been used.\n        if len(visited_wires.intersection(wires_set)) > 0:\n            raise DeviceError(\n                f\"Operation {current_gate.name} cannot be used after other Operation applied in the same qubit \"\n            )\n        input_wires.append(current_gate.wires)\n        input_vectors.append(current_gate.parameters[0])\n        input_batch_size.append(current_gate.batch_size)\n        list_copy.pop(0)\n        visited_wires = visited_wires.union(wires_set)\n\n    if len(input_wires) > 0:\n        final_wires = input_wires[0]\n        final_vector = input_vectors[0]\n        final_batch_size = input_batch_size[0]\n\n        # Merge all parameters and qubits into a single one.\n        for w, v, b in zip(input_wires[1:], input_vectors[1:], input_batch_size[1:]):\n            final_vector = final_vector[..., :, None] * v[..., None, :]\n            final_batch_size = final_batch_size or b\n            final_wires = final_wires + w\n\n            if final_batch_size:\n                final_vector = reshape(final_vector, (final_batch_size, -1))\n            else:\n                final_vector = flatten(final_vector)\n\n        AmplitudeEmbedding(final_vector, wires=final_wires)\n\n    new_operations = []\n    for gate in not_amplitude_embedding:\n        new_operations.append(gate)\n\n    new_tape = type(tape)(new_operations, tape.measurements, shots=tape.shots)\n\n    def null_postprocessing(results):\n        \"\"\"A postprocesing function returned by a transform that only converts the batch of results\n        into a result for a single ``QuantumTape``.\n        \"\"\"\n        return results[0]\n\n    return [new_tape], null_postprocessing\n", "path": "pennylane/transforms/optimization/merge_amplitude_embedding.py"}], "after_files": [{"content": "# Copyright 2018-2021 Xanadu Quantum Technologies Inc.\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Transform for merging AmplitudeEmbedding gates in a quantum circuit.\"\"\"\nfrom typing import Sequence, Callable\n\nfrom pennylane.transforms import transform\nfrom pennylane.tape import QuantumTape\nfrom pennylane import AmplitudeEmbedding\nfrom pennylane._device import DeviceError\nfrom pennylane.math import flatten, reshape\nfrom pennylane.queuing import QueuingManager\n\n\n@transform\ndef merge_amplitude_embedding(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):\n    r\"\"\"Quantum function transform to combine amplitude embedding templates that act on different qubits.\n\n    Args:\n        tape (QNode or QuantumTape or Callable): A quantum circuit.\n\n    Returns:\n        qnode (QNode) or quantum function (Callable) or tuple[List[.QuantumTape], function]: The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.\n\n\n    **Example**\n\n    >>> dev = qml.device('default.qubit', wires=4)\n\n    You can apply the transform directly on :class:`QNode`:\n\n    .. code-block:: python\n\n        @qml.transforms.merge_amplitude_embedding\n        @qml.qnode(device=dev)\n        def circuit():\n            qml.CNOT(wires = [0,1])\n            qml.AmplitudeEmbedding([0,1], wires = 2)\n            qml.AmplitudeEmbedding([0,1], wires = 3)\n            return qml.state()\n\n    >>> circuit()\n    [1.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j]\n\n    .. details::\n        :title: Usage Details\n\n        You can also apply it on quantum function.\n\n        .. code-block:: python\n\n            def qfunc():\n                qml.CNOT(wires = [0,1])\n                qml.AmplitudeEmbedding([0,1], wires = 2)\n                qml.AmplitudeEmbedding([0,1], wires = 3)\n                return qml.state()\n\n        The circuit before compilation will not work because of using two amplitude embedding.\n\n        Using the transformation we can join the different amplitude embedding into a single one:\n\n        >>> optimized_qfunc = qml.transforms.merge_amplitude_embedding(qfunc)\n        >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)\n        >>> print(qml.draw(optimized_qnode)())\n        0: \u2500\u256d\u25cf\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2524  State\n        1: \u2500\u2570X\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2524  State\n        2: \u2500\u256dAmplitudeEmbedding(M0)\u2500\u2524  State\n        3: \u2500\u2570AmplitudeEmbedding(M0)\u2500\u2524  State\n        M0 =\n        [0.+0.j 0.+0.j 0.+0.j 1.+0.j]\n\n    \"\"\"\n    new_operations = []\n    visited_wires = set()\n    input_wires, input_vectors, input_batch_size = [], [], []\n    for current_gate in tape.operations:\n        wires_set = set(current_gate.wires)\n\n        # Check if the current gate is an AmplitudeEmbedding.\n        if not isinstance(current_gate, AmplitudeEmbedding):\n            new_operations.append(current_gate)\n            visited_wires = visited_wires.union(wires_set)\n            continue\n\n        # Check the qubits have not been used.\n        if len(visited_wires.intersection(wires_set)) > 0:\n            raise DeviceError(\n                f\"Operation {current_gate.name} cannot be used after other Operation applied in the same qubit \"\n            )\n        input_wires.append(current_gate.wires)\n        input_vectors.append(current_gate.parameters[0])\n        input_batch_size.append(current_gate.batch_size)\n        visited_wires = visited_wires.union(wires_set)\n\n    if len(input_wires) > 0:\n        final_wires = input_wires[0]\n        final_vector = input_vectors[0]\n        final_batch_size = input_batch_size[0]\n\n        # Merge all parameters and qubits into a single one.\n        for w, v, b in zip(input_wires[1:], input_vectors[1:], input_batch_size[1:]):\n            final_vector = final_vector[..., :, None] * v[..., None, :]\n            final_batch_size = final_batch_size or b\n            final_wires = final_wires + w\n\n            if final_batch_size:\n                final_vector = reshape(final_vector, (final_batch_size, -1))\n            else:\n                final_vector = flatten(final_vector)\n\n        with QueuingManager.stop_recording():\n            new_operations.insert(0, AmplitudeEmbedding(final_vector, wires=final_wires))\n\n    new_tape = type(tape)(new_operations, tape.measurements, shots=tape.shots)\n\n    def null_postprocessing(results):\n        \"\"\"A postprocesing function returned by a transform that only converts the batch of results\n        into a result for a single ``QuantumTape``.\n        \"\"\"\n        return results[0]\n\n    return [new_tape], null_postprocessing\n", "path": "pennylane/transforms/optimization/merge_amplitude_embedding.py"}]}