Repository: PennyLaneAI/pennylane

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.


Discussion:
Thanks for opening this issue @akatief .  This does indeed look to be a bug. We'll look into this and get back to you.

Resolving patch:
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 @@ def qfunc():
         [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 @@ def qfunc():
         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 @@ def qfunc():
             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)
 


Test patch:
diff --git a/tests/transforms/test_optimization/test_merge_amplitude_embedding.py b/tests/transforms/test_optimization/test_merge_amplitude_embedding.py
--- a/tests/transforms/test_optimization/test_merge_amplitude_embedding.py
+++ b/tests/transforms/test_optimization/test_merge_amplitude_embedding.py
@@ -55,10 +55,10 @@ def circuit():
             qml.AmplitudeEmbedding([0.0, 1.0], wires=0)
             qml.AmplitudeEmbedding([0.0, 1.0], wires=1)
             qml.Hadamard(wires=0)
-            qml.Hadamard(wires=0)
+            qml.Hadamard(wires=1)
             return qml.state()
 
-        assert np.allclose(circuit()[0], 1)
+        assert qml.math.allclose(circuit(), np.array([1, -1, -1, 1]) / 2)
 
     def test_repeated_qubit(self):
         """Check that AmplitudeEmbedding cannot be applied if the qubit has already been used."""
@@ -93,8 +93,8 @@ def test_broadcasting(self):
         """Test that merging preserves the batch dimension"""
         dev = qml.device("default.qubit", wires=3)
 
-        @qml.qnode(dev)
         @qml.transforms.merge_amplitude_embedding
+        @qml.qnode(dev)
         def qnode():
             qml.AmplitudeEmbedding([[1, 0], [0, 1]], wires=0)
             qml.AmplitudeEmbedding([1, 0], wires=1)