Repository: PennyLaneAI/pennylane
Pull request #7426: [Capture] Fix plxpr transforms that don't take dynamic wires into account
Author: login: mudit2812
Created: 2025-05-13T21:18:14Z | Merged: 2025-05-30T19:22:31Z

[sc-90121]

**Context:**
Fixes #7349.

**Description of the Change:**
* Update plxpr `cancel_inverses` to interpret all previous ops if an op with dynamic wires is encountered.
* Update plxpr `merge_rotations` to interpret all previous ops if an op with dynamic wires is encountered.
* Update plxpr `single_qubit_fusion` to interpret all previous ops if an op with dynamic wires is encountered.
* Update plxpr `merge_amplitude_embedding` to merge all `AmplitudeEmbedding`s found so far and all ops encountered so far if an op with dynamic wires is encountered.
  * If an `AmplitudeEmbedding` is found after a non-`AmplitudeEmbedding` op _with dynamic wires_, an error is raised.
  * If an `AmplitudeEmbedding` _with dynamic wires_ is found after any non-`AmplitudeEmbedding` op, an error is raised.
  * If the above conditions _are not_ met, but there are multiple `AmplitudeEmbedding`s, no error is raised regardless of presence of dynamic wires.
* Update program capture sharp bits page to notify users that non-native transforms may have incorrect behaviour with dynamic wires. 

**Benefits:**
* Correct behaviour of above transforms with program capture
* Better information to users about sharp bits

**Possible Drawbacks:**

**Related GitHub Issues:**
#7349 

Linked issue #7349: [BUG] Some optimization transforms do not work correctly with dynamic wires and program capture
### Expected behavior

In some optimization transforms, the existence of operators on dynamic wires does not correctly stop the optimizations from occurring, leading to incorrect circuits. The transforms affected by this are:
* `cancel_inverses`
* `merge_rotations`
* `single_qubit_fusion`
* `merge_amplitude_embedding`

### Actual behavior

Operators that shouldn't be able to get optimized get optimized, leading to incorrect circuits. An example using `cancel_inverses` and `merge_rotations` is shown below.

### Additional information

_No response_

### Source code

```shell
# Cancel inverses:

import pennylane as qml
import jax

qml.capture.enable()

@qml.capture.expand_plxpr_transforms
@qml.transforms.cancel_inverses
def f(x):
    qml.X(0)
    qml.T(x)
    qml.X(0)
    return qml.expval(qml.Z(0))

print("Cancel inverses:")
print(jax.make_jaxpr(f)(0))


# Merge Rotations
@qml.capture.expand_plxpr_transforms
@qml.transforms.merge_rotations
def g(x):
    qml.RX(1.5, 0)
    qml.RZ(2.5, x)
    qml.RX(3.5, 0)
    return qml.expval(qml.Z(0))

print("\nMerge rotations:")
print(jax.make_jaxpr(g)(0))
```

### Tracebacks

```shell
Cancel inverses:
{ lambda ; a:i64[]. let
    _:AbstractOperator() = T[n_wires=1] a
    b:AbstractOperator() = PauliZ[n_wires=1] 0
    c:AbstractMeasurement(n_wires=None) = expval_obs b
  in (c,) }

Merge rotations:
{ lambda ; a:i64[]. let
    _:AbstractOperator() = RX[n_wires=1] 5.0 0
    _:AbstractOperator() = RZ[n_wires=1] 2.5 a
    b:AbstractOperator() = PauliZ[n_wires=1] 0
    c:AbstractMeasurement(n_wires=None) = expval_obs b
  in (c,) }
```

### System information

```shell
PL dev
```

### Existing GitHub issues

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

Diff:
diff --git a/doc/news/program_capture_sharp_bits.rst b/doc/news/program_capture_sharp_bits.rst
index aa96faa1c6f..90baf81aced 100644
--- a/doc/news/program_capture_sharp_bits.rst
+++ b/doc/news/program_capture_sharp_bits.rst
@@ -448,6 +448,8 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit whose wires are dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 92d2e6b73f4..b7eefe70758 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -585,6 +585,10 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Documentation 📝</h3>
 
+* The entry in the :doc:`/news/program_capture_sharp_bits` page for transforms has been updated; non-native transforms being applied
+  to QNodes wherein operators have dynamic wires can lead to incorrect results.
+  [(#7426)](https://github.com/PennyLaneAI/pennylane/pull/7426)
+
 * Fixed the wrong `theta` to `phi` in :class:`~pennylane.IsingXY`.
  [(#7427)](https://github.com/PennyLaneAI/pennylane/pull/7427)
 
@@ -600,6 +604,14 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixed a bug where certain transforms with a native program capture implementation give incorrect results when
+  dynamic wires were present in the circuit. The affected transforms were:
+  * :func:`~pennylane.transforms.cancel_inverses`
+  * :func:`~pennylane.transforms.merge_rotations`
+  * :func:`~pennylane.transforms.single_qubit_fusion`
+  * :func:`~pennylane.transforms.merge_amplitude_embedding`
+  [(#7426)](https://github.com/PennyLaneAI/pennylane/pull/7426)
+
 * The `Operator.pow` method has been fixed to raise to the power of 2 the qutrit operators `~.TShift`, `~.TClock`, and `~.TAdd`.
   [(#7505)](https://github.com/PennyLaneAI/pennylane/pull/7505)
 
diff --git a/pennylane/transforms/optimization/cancel_inverses.py b/pennylane/transforms/optimization/cancel_inverses.py
index 68dc89111e0..7325d3c86fa 100644
--- a/pennylane/transforms/optimization/cancel_inverses.py
+++ b/pennylane/transforms/optimization/cancel_inverses.py
@@ -15,6 +15,7 @@
 
 from functools import lru_cache, partial
 
+from pennylane.math import is_abstract
 from pennylane.ops.op_math import Adjoint
 from pennylane.ops.qubit.attributes import (
     self_inverses,
@@ -120,7 +121,17 @@ def interpret_operation(self, op: Operator):
                 return super().interpret_operation(op)
 
             prev_op = self.previous_ops.get(op.wires[0], None)
-            if prev_op is None:
+            dyn_wires = {w for w in op.wires if is_abstract(w)}
+            other_saved_wires = set(self.previous_ops.keys()) - dyn_wires
+
+            if prev_op is None or (dyn_wires and other_saved_wires):
+                # If there are dynamic wires, we need to make sure that there are no
+                # other wires in `self.previous_ops`, otherwise we can't cancel. If
+                # there are other wires but no other op on the same dynamic wire(s),
+                # there isn't anything to cancel, so we just add the current op to
+                # `self.previous_ops` and return.
+                if dyn_wires and (prev_op is None or other_saved_wires):
+                    self.interpret_all_previous_ops()
                 for w in op.wires:
                     self.previous_ops[w] = op
                 return []
diff --git a/pennylane/transforms/optimization/merge_amplitude_embedding.py b/pennylane/transforms/optimization/merge_amplitude_embedding.py
index 4e29822bac7..af129ba27ee 100644
--- a/pennylane/transforms/optimization/merge_amplitude_embedding.py
+++ b/pennylane/transforms/optimization/merge_amplitude_embedding.py
@@ -20,10 +20,10 @@
 import pennylane as qml
 from pennylane import AmplitudeEmbedding
 from pennylane.exceptions import DeviceError
-from pennylane.math import flatten, reshape
+from pennylane.math import flatten, is_abstract, reshape
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms import transform
+from pennylane.transforms.core import TransformError, transform
 from pennylane.typing import PostprocessingFn
 
 
@@ -48,8 +48,13 @@ class MergeAmplitudeEmbeddingInterpreter(PlxprInterpreter):
 
         def __init__(self):
             self._env = {}
+            self.dynamic_wires_encountered = False
             self.previous_ops = []
-            self.state = {"visited_wires": set()}
+            # * visited_wires (set): tracks all wires we have encountered so far.
+            # * dynamic_wires_found (bool): True if we have encountered any non-AmplitudeEmbedding
+            #   ops that have dynamic wires so far.
+            # * ops_found (bool): True if we have encountered any non-AmplitudeEmbedding ops so far.
+            self.state = {"visited_wires": set(), "dynamic_wires_found": False, "ops_found": False}
             self.input_wires, self.input_vectors, self.input_batch_size = [], [], []
 
         def setup(self) -> None:
@@ -57,6 +62,10 @@ def setup(self) -> None:
             self.previous_ops = []
             self.input_wires, self.input_vectors, self.input_batch_size = [], [], []
 
+        def cleanup(self) -> None:
+            """Clean up the interpreter after evaluation."""
+            self.state = {"visited_wires": set(), "dynamic_wires_found": False, "ops_found": False}
+
         def interpret_operation(self, op: Operator) -> None:
             """Interpret a PennyLane operation instance.
 
@@ -76,13 +85,33 @@ def interpret_operation(self, op: Operator) -> None:
             so the output will not affect later equations in the circuit.
 
             """
+
             if not isinstance(op, AmplitudeEmbedding):
+                if any(is_abstract(w) for w in op.wires):
+                    if self.input_wires:
+                        self._merge_and_insert_at_the_start()
+                    self.interpret_all_previous_ops()
+                    self.state["dynamic_wires_found"] = True
+
+                self.state["ops_found"] = True
                 self.previous_ops.append(op)
                 self.state["visited_wires"] = self.state["visited_wires"].union(set(op.wires))
                 return
 
+            if self.state["dynamic_wires_found"]:
+                raise TransformError(
+                    "Cannot apply qml.AmplitudeEmbedding after operators with dynamic wires as it "
+                    "is indeterminable if the wires overlap."
+                )
+
+            if self.state["ops_found"] and any(is_abstract(w) for w in op.wires):
+                raise TransformError(
+                    "Cannot apply qml.AmplitudeEmbedding with dynamic wires after other operators "
+                    "as it is indeterminable if the wires overlap."
+                )
+
             if len(self.state["visited_wires"].intersection(set(op.wires))) > 0:
-                raise DeviceError(
+                raise TransformError(
                     "qml.AmplitudeEmbedding cannot be applied on wires already used by other operations."
                 )
 
@@ -112,10 +141,8 @@ def _merge_and_insert_at_the_start(self) -> None:
                 else:
                     final_vector = flatten(final_vector)
 
-            # pylint: disable=protected-access
-            self.previous_ops.insert(
-                0, qml.AmplitudeEmbedding._primitive.impl(final_vector, wires=final_wires)
-            )
+            with qml.capture.pause():
+                self.previous_ops.insert(0, qml.AmplitudeEmbedding(final_vector, wires=final_wires))
             # Clear history of amplitude embedding gates since we've merged
             self.input_wires, self.input_vectors, self.input_batch_size = [], [], []
 
@@ -208,10 +235,19 @@ def _(self, *invals, jaxpr_branches, consts_slices, args_slice):
         new_consts_slices = []
         end_const_ind = len(jaxpr_branches)
 
-        # Store seen wires before we begin to process the branches
-        # (create copies as to not accidently mutate the original state)
+        # Store state before we begin to process the branches
+        # (create copies as to not accidently mutate the original state).
+        # We cannot just copy self.state because a shallow copy would not
+        # create a copy of `visited_wires`, which is a set.
+        # We cannot use deepcopy as `visited_wires` may have tracers inside,
+        # which have hashes specific to the instance. Copying these will cause
+        # the dynamic wires in the original and copy to be different.
         initial_wires = copy(self.state["visited_wires"])
-        visited_wires = copy(initial_wires)
+        curr_wires = copy(self.state["visited_wires"])
+        initial_dynamic_wires_found = self.state["dynamic_wires_found"]
+        curr_dynamic_wires_found = self.state["dynamic_wires_found"]
+        initial_ops_found = self.state["ops_found"]
+        curr_ops_found = self.state["ops_found"]
 
         for const_slice, jaxpr in zip(consts_slices, jaxpr_branches):
             consts = invals[const_slice]
@@ -221,14 +257,20 @@ def _(self, *invals, jaxpr_branches, consts_slices, args_slice):
             else:
                 new_jaxpr = jaxpr_to_jaxpr(copy(self), jaxpr, consts, *args)
 
-                # Update wires we've seen so far so collisions with
-                # newly seen wires from the branches continue to be
+                # Update state so far so collisions with
+                # newly seen states from the branches continue to be
                 # detected after the cond
-                visited_wires |= self.state["visited_wires"]
-
-                # Reset visited wires for the next branch so we don't get false positive collisions
-                # (copy so if state mutates we preserved true initial wires)
-                self.state["visited_wires"] = copy(initial_wires)
+                curr_wires |= self.state["visited_wires"]
+                curr_dynamic_wires_found = self.state["dynamic_wires_found"]
+                curr_ops_found = self.state["ops_found"]
+
+                # Reset state for the next branch so we don't get false positive collisions
+                # (copy so if state mutates we preserved true initial state)
+                self.state = {
+                    "visited_wires": copy(initial_wires),
+                    "dynamic_wires_found": initial_dynamic_wires_found,
+                    "ops_found": initial_ops_found,
+                }
 
                 new_jaxprs.append(new_jaxpr.jaxpr)
                 new_consts.extend(new_jaxpr.consts)
@@ -237,8 +279,12 @@ def _(self, *invals, jaxpr_branches, consts_slices, args_slice):
                 )
                 end_const_ind += len(new_jaxpr.consts)
 
-        # Reset visited wires to all wires encountered in the cond
-        self.state["visited_wires"] = visited_wires
+        # Reset state to all updates from all branches in the cond
+        self.state = {
+            "visited_wires": curr_wires,
+            "dynamic_wires_found": curr_dynamic_wires_found,
+            "ops_found": curr_ops_found,
+        }
 
         new_args_slice = slice(end_const_ind, None)
         return cond_prim.bind(
@@ -255,6 +301,9 @@ def _(self, *invals, **params):
         # Make sure to record that we have visited the wires on this measurement
         # in order to be able to detect potential wire collisions with future AE gates
         self.state["visited_wires"] = self.state["visited_wires"].union(set(invals))
+        self.state["dynamic_wires_found"] = any(is_abstract(w) for w in invals)
+        self.state["ops_found"] = True
+
         # pylint: disable=protected-access
         if len(self.input_wires) > 0:
             self._merge_and_insert_at_the_start()
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 048d7f339ba..39ee815f9dc 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -78,7 +78,7 @@ def _interpret_previous_ops_on_wires(self, wires) -> None:
             for prev_op in previous_ops_on_wires:
                 super().interpret_operation(prev_op)
 
-        # pylint: disable=inconsistent-return-statements
+        # pylint: disable=inconsistent-return-statements,too-many-branches
         def interpret_operation(self, op: Operator):
             """Interpret a PennyLane operation instance.
 
@@ -101,14 +101,23 @@ def interpret_operation(self, op: Operator):
                 return self._update_previous_ops(op)
 
             previous_op = self.previous_ops.get(op.wires[0])
-            if previous_op is None:
+            dyn_wires = {w for w in op.wires if qml.math.is_abstract(w)}
+            other_saved_wires = set(self.previous_ops.keys()) - dyn_wires
+            if previous_op is None or (dyn_wires and other_saved_wires):
+                # If there are dynamic wires, we need to make sure that there are no
+                # other wires in `self.previous_ops`, otherwise we can't merge. If
+                # there are other wires but no other op on the same dynamic wire(s),
+                # there isn't anything to merge, so we just add the current op to
+                # `self.previous_ops` and return.
+                if dyn_wires and (previous_op is None or other_saved_wires):
+                    self._interpret_remaining_ops()
                 for w in op.wires:
                     self.previous_ops[w] = op
                 return
 
             # pylint: disable = unidiomatic-typecheck
             # Can't use `isinstance` since op could be a subclass of type(previous_op)
-            can_merge = (op.wires == previous_op.wires) and (type(op) == type(previous_op))
+            can_merge = op.wires == previous_op.wires and type(op) == type(previous_op)
             if not can_merge:
                 self._interpret_previous_ops_on_wires(op.wires)
                 return self._update_previous_ops(op)
@@ -136,6 +145,12 @@ def interpret_operation(self, op: Operator):
                 or qml.math.requires_grad(cumulative_angles)
                 or not angles_cancel
             )
+
+            if any(qml.math.is_abstract(w) for w in op.wires):
+                for w in op.wires:
+                    del self.previous_ops[w]
+                self._interpret_remaining_ops()
+
             if keep_merged_op:
                 # pylint: disable = protected-access
                 new_op = op._primitive.impl(*cumulative_angles, wires=op.wires)
diff --git a/pennylane/transforms/optimization/single_qubit_fusion.py b/pennylane/transforms/optimization/single_qubit_fusion.py
index 3459002e611..0cabf94f8fa 100644
--- a/pennylane/transforms/optimization/single_qubit_fusion.py
+++ b/pennylane/transforms/optimization/single_qubit_fusion.py
@@ -89,13 +89,20 @@ def _handle_non_fusible_op(self, op: Operator) -> list:
             """Handle an operation that cannot be fused into a Rot gate."""
 
             previous_ops_on_wires = self._retrieve_prev_ops_same_wire(op)
+            dyn_wires = {w for w in op.wires if qml.math.is_abstract(w)}
+            other_saved_wires = set(self.previous_ops.keys()) - dyn_wires
+
+            if dyn_wires and other_saved_wires:
+                # We cannot guarantee that ops on other static or dynamic wires won't have wire
+                # overlap with the current op, so we need to interpret all of them.
+                self.interpret_all_previous_ops()
 
             res = []
             for prev_op in previous_ops_on_wires:
-                # pylint: disable=protected-access
-                rot = qml.Rot._primitive.impl(
-                    *qml.math.stack(prev_op.single_qubit_rot_angles()), wires=prev_op.wires
-                )
+                with qml.capture.pause():
+                    rot = qml.Rot(
+                        *qml.math.stack(prev_op.single_qubit_rot_angles()), wires=prev_op.wires
+                    )
                 res.append(super().interpret_operation(rot))
 
             res.append(super().interpret_operation(op))
@@ -108,9 +115,20 @@ def _handle_fusible_op(self, op: Operator, cumulative_angles: TensorLike) -> lis
             # Only single-qubit gates are considered for fusion
             op_wire = op.wires[0]
 
-            prev_op = self.previous_ops.get(op_wire)
-            if prev_op is None:
-                self.previous_ops[op_wire] = op
+            prev_op = self.previous_ops.get(op.wires[0], None)
+            dyn_wires = {w for w in op.wires if qml.math.is_abstract(w)}
+            other_saved_wires = set(self.previous_ops.keys()) - dyn_wires
+
+            if prev_op is None or (dyn_wires and other_saved_wires):
+                # If there are dynamic wires, we need to make sure that there are no
+                # other wires in `self.previous_ops`, otherwise we can't fuse. If
+                # there are other wires but no other op on the same dynamic wire(s),
+                # there isn't anything to fuse, so we just add the current op to
+                # `self.previous_ops` and return.
+                if dyn_wires and (prev_op is None or other_saved_wires):
+                    self.interpret_all_previous_ops()
+                for w in op.wires:
+                    self.previous_ops[w] = op
                 return []
 
             prev_op_angles = qml.math.stack(prev_op.single_qubit_rot_angles())
@@ -128,8 +146,8 @@ def _handle_fusible_op(self, op: Operator, cumulative_angles: TensorLike) -> lis
                     rtol=0,
                 )
             ):
-                # pylint: disable=protected-access
-                new_rot = qml.Rot._primitive.impl(*cumulative_angles, wires=op.wires)
+                with qml.capture.pause():
+                    new_rot = qml.Rot(*cumulative_angles, wires=op.wires)
                 self.previous_ops[op_wire] = new_rot
             else:
                 del self.previous_ops[op_wire]
diff --git a/tests/capture/transforms/test_capture_cancel_inverses.py b/tests/capture/transforms/test_capture_cancel_inverses.py
index c2b13334b9f..d00d3cd8803 100644
--- a/tests/capture/transforms/test_capture_cancel_inverses.py
+++ b/tests/capture/transforms/test_capture_cancel_inverses.py
@@ -169,15 +169,109 @@ def f():
         ops = collector.state["ops"]
         meas = collector.state["measurements"]
 
-        expected_ops = [
-            qml.Identity(),
-            qml.Identity(),
-        ]
+        expected_ops = [qml.Identity(), qml.Identity()]
         assert ops == expected_ops
 
         expected_meas = [qml.expval(qml.Z(0))]
         assert meas == expected_meas
 
+    def test_dynamic_wires_between_static_wires(self):
+        """Test that operations with dynamic wires between operations with static
+        wires cause cancellation to not happen."""
+
+        @CancelInversesInterpreter()
+        def f(w):
+            qml.H(0)
+            qml.T(w)
+            qml.H(0)
+            qml.adjoint(qml.T(w))
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(0)
+        assert len(jaxpr.eqns) == 7
+
+        dyn_wire = 0
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+        expected_meas = [qml.expval(qml.Z(0))]
+
+        expected_ops = [qml.H(0), qml.T(0), qml.H(0), qml.adjoint(qml.T(0))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+        dyn_wire = 1
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.H(0), qml.T(1), qml.H(0), qml.adjoint(qml.T(1))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+    def test_same_dyn_wires_cancel(self):
+        """Test that ops on the same dynamic wires get cancelled."""
+
+        @CancelInversesInterpreter()
+        def f(w):
+            qml.H(0)
+            qml.T(w)
+            qml.adjoint(qml.T(w))
+            qml.H(0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(0)
+        assert len(jaxpr.eqns) == 4
+
+        dyn_wire = 0
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.H(0), qml.H(0)]
+        expected_meas = [qml.expval(qml.Z(0))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+    def test_different_dyn_wires_interleaved(self):
+        """Test that ops on different dynamic wires interleaved with each other
+        do not cancel."""
+
+        @CancelInversesInterpreter()
+        def f(w1, w2):
+            qml.H(w1)
+            qml.X(w2)
+            qml.H(w1)
+            qml.X(w2)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(0, 0)
+        assert len(jaxpr.eqns) == 6
+
+        dyn_wires = (0, 0)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, *dyn_wires)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.H(0), qml.X(0), qml.H(0), qml.X(0)]
+        expected_meas = [qml.expval(qml.Z(0))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+        dyn_wires = (0, 1)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, *dyn_wires)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.H(0), qml.X(1), qml.H(0), qml.X(1)]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
     def test_ctrl_higher_order_primitive(self):
         """Test that ctrl higher order primitives are transformed correctly."""
 
diff --git a/tests/capture/transforms/test_capture_merge_amplitude_embedding.py b/tests/capture/transforms/test_capture_merge_amplitude_embedding.py
index 87eb559b5fc..691d9363efd 100644
--- a/tests/capture/transforms/test_capture_merge_amplitude_embedding.py
+++ b/tests/capture/transforms/test_capture_merge_amplitude_embedding.py
@@ -16,9 +16,10 @@
 import pytest
 
 import pennylane as qml
-from pennylane.exceptions import DeviceError
+from pennylane.transforms.core import TransformError
 
 jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
 
 
 def _find_eq_with_name(jaxpr, name):
@@ -48,8 +49,8 @@ def _find_eq_with_name(jaxpr, name):
 pytestmark = [pytest.mark.jax, pytest.mark.usefixtures("enable_disable_plxpr")]
 
 
-class TestRepeatedQubitDeviceErrors:
-    """Test DeviceError is raised when operations exist before the AmplitudeEmbedding operators."""
+class TestRepeatedQubitTransformErrors:
+    """Test TransformError is raised when operations exist before the AmplitudeEmbedding operators."""
 
     def test_repeated_wire_with_mcm(self):
         """Test that an error is raised if AmplitudeEmbedding acts on the same wires as an MCM."""
@@ -57,13 +58,27 @@ def test_repeated_wire_with_mcm(self):
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc():
             qml.Hadamard(wires=2)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.measure(wires=1)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
             return qml.expval(qml.Z(0))
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
+            match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
+        ):
+            jax.make_jaxpr(qfunc)()
+
+    def test_repeated_wire_error(self):
+        """Test that an error is raised if AmplitudeEmbedding acts on the same wire"""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def qfunc():
+            qml.X(0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+
+        with pytest.raises(
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             jax.make_jaxpr(qfunc)()
@@ -72,13 +87,13 @@ def test_repeated_traced_wire_error(self):
         """Test that an error is raised if AmplitudeEmbedding acts on the same traced wire"""
 
         @MergeAmplitudeEmbeddingInterpreter()
-        def qfunc(wire):
-            qml.X(wire)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=wire)
+        def qfunc(w):
+            qml.X(w)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=w)
 
         with pytest.raises(
-            DeviceError,
-            match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
+            TransformError,
+            match="Cannot apply qml.AmplitudeEmbedding after operators with dynamic wires.",
         ):
             jax.make_jaxpr(qfunc)(1)
 
@@ -88,11 +103,11 @@ def test_simple_repeated_qubit_error(self):
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc():
             qml.CNOT(wires=[0.0, 1.0])
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             jax.make_jaxpr(qfunc)()
@@ -101,8 +116,8 @@ def test_repeated_qubit_error_before_higher_order_prim(self):
         """Test that wire collision can be detected before a higher-order primitive is applied."""
 
         def ctrl_fn():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
         @MergeAmplitudeEmbeddingInterpreter()
         def f():
@@ -110,7 +125,7 @@ def f():
             qml.ctrl(ctrl_fn, [2, 3])()
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             jax.make_jaxpr(f)()
@@ -124,11 +139,11 @@ def ctrl_fn():
         @MergeAmplitudeEmbeddingInterpreter()
         def f():
             qml.ctrl(ctrl_fn, [2])()  # ctrl_fn has wires 0
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             jax.make_jaxpr(f)()
@@ -141,21 +156,21 @@ def test_collision_before_cond_prim(self, collision_wire):
         def f(x):
             @qml.cond(x > 2)
             def cond_f():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
 
             @cond_f.else_if(x > 1)
             def _():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
             @cond_f.otherwise
             def _():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=2)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=2)
 
             qml.X(collision_wire)  # Each branch of the cond will have a collision with this wire
             cond_f()
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             jax.make_jaxpr(f)(3)
@@ -179,10 +194,10 @@ def _():
                 qml.X(2)
 
             cond_f()  # Each branch of the cond will have a collision with this wire
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=collision_wire)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=collision_wire)
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             jax.make_jaxpr(f)(3)
@@ -208,10 +223,10 @@ def _():
             qml.X(0)
             cond_f()
             # visited wires after cond should be 0 (before cond) and 1, 2, 3 (during cond)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=collision_wire)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=collision_wire)
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             jax.make_jaxpr(f)(3)
@@ -227,24 +242,24 @@ def loop(i):
 
             @qml.cond(x > 2)
             def cond_f():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
             @cond_f.else_if(x > 1)
             def _():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
             @cond_f.otherwise
             def _():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
             loop()  # Contains wires 0
             cond_f()  # Also contains wires 0
 
         with pytest.raises(
-            DeviceError,
+            TransformError,
             match="qml.AmplitudeEmbedding cannot be applied on wires already used by other operations.",
         ):
             args = (3,)
@@ -264,13 +279,13 @@ def qfunc(state1, state2):
             qml.AmplitudeEmbedding(state2, wires=1)
             return qml.expval(qml.Z(0))
 
-        args = (jax.numpy.array([0.0, 1.0]), jax.numpy.array([0.0, 1.0]))
+        args = (jnp.array([0.0, 1.0]), jnp.array([0.0, 1.0]))
         jaxpr = jax.make_jaxpr(qfunc)(*args)
         collector = CollectOpsandMeas()
         collector.eval(jaxpr.jaxpr, jaxpr.consts, *args)
 
         expected_ops = [
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 0.0, 0.0, 1.0]), wires=[0, 1]),
+            qml.AmplitudeEmbedding(jnp.array([0.0, 0.0, 0.0, 1.0]), wires=[0, 1]),
             qml.Hadamard(wires=[0]),
         ]
 
@@ -288,9 +303,9 @@ def test_circuit_with_traced_wires(self):
 
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc(wires1, wires2):
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=wires1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=wires1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=wires2)
             qml.Hadamard(wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=wires2)
             return qml.expval(qml.Z(0))
 
         args = (0, 1)
@@ -318,7 +333,7 @@ def test_circuit_with_no_merge_required(self):
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc():
             qml.Hadamard(wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
             return qml.expval(qml.Z(0))
 
         jaxpr = jax.make_jaxpr(qfunc)()
@@ -345,10 +360,10 @@ def test_merge_simple(self):
 
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.Hadamard(wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 0.0, 0.0, 1.0]), wires=[2, 3])
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 0.0, 0.0, 1.0]), wires=[2, 3])
             return qml.expval(qml.Z(0))
 
         jaxpr = jax.make_jaxpr(qfunc)()
@@ -376,11 +391,11 @@ def test_batch_preservation(self):
 
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc():
-            qml.AmplitudeEmbedding(jax.numpy.array([[0, 1], [1, 0]]), wires=0)  # |1> and |0>
+            qml.AmplitudeEmbedding(jnp.array([[0, 1], [1, 0]]), wires=0)  # |1> and |0>
             qml.AmplitudeEmbedding(
-                jax.numpy.array([1, 0]), wires=1
+                jnp.array([1, 0]), wires=1
             )  # |0> (batch will be extended to |0> and |0>)
-            qml.AmplitudeEmbedding(jax.numpy.array([[0, 1], [1, 0]]), wires=2)  # |1> and |0>
+            qml.AmplitudeEmbedding(jnp.array([[0, 1], [1, 0]]), wires=2)  # |1> and |0>
             return qml.expval(qml.Z(0))
 
         jaxpr = jax.make_jaxpr(qfunc)()
@@ -407,8 +422,8 @@ def test_returned_ops_are_not_merged(self):
 
         @MergeAmplitudeEmbeddingInterpreter()
         def f():
-            qml.AmplitudeEmbedding(jax.numpy.array([1, 0]), wires=0)
-            return qml.AmplitudeEmbedding(jax.numpy.array([1, 0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([1, 0]), wires=0)
+            return qml.AmplitudeEmbedding(jnp.array([1, 0]), wires=1)
 
         jaxpr = jax.make_jaxpr(f)()
         collector = CollectOpsandMeas()
@@ -422,14 +437,86 @@ def f():
         ops = collector.state["ops"]
         assert ops == expected_ops
 
+    def test_amplitude_embedding_before_dynamic_wires_op(self):
+        """Test that AmplitudeEmbeddings are merged and applied
+        before ops with dynamic wires."""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(w):
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=1)
+            qml.H(w)
+
+        jaxpr = jax.make_jaxpr(circuit)(0)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 0)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.AmplitudeEmbedding(qml.math.array([1.0, 0.0, 0.0, 0.0]), wires=[0, 1]),
+            qml.H(0),
+        ]
+        assert ops == expected_ops
+
+    def test_dynamic_wire_ops_before_error(self):
+        """Test that an error is raised if ops with dynamic wires are used before AmplitudeEmbedding"""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def f(w):
+            qml.H(w)
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=[0])
+
+        with pytest.raises(
+            TransformError,
+            match="Cannot apply qml.AmplitudeEmbedding after operators with dynamic wires.",
+        ):
+            _ = jax.make_jaxpr(f)(2)
+
+    def test_dynamic_wire_embeddings(self):
+        """Test that AmplitudeEmbeddings with dynamic wires can be merged. If there is overlap,
+        an error will be raised at runtime."""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def f(w):
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0]), [w])
+            qml.H(0)
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), [1])
+
+        jaxpr = jax.make_jaxpr(f)(0)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 0)
+        ops = collector.state["ops"]
+        expected = [qml.AmplitudeEmbedding(qml.math.array([1.0, 0.0, 0.0, 0.0]), [0, 1]), qml.H(0)]
+
+        assert ops == expected
+
+        # If first embedding has wire 1, error will happen at runtime
+        collector = CollectOpsandMeas()
+        with pytest.raises(qml.wires.WireError, match="Wires must be unique"):
+            collector.eval(jaxpr.jaxpr, jaxpr.consts, 1)
+
+    def test_dynamic_wire_embedding_after_op(self):
+        """Test that an error is raised if an AmplitudeEmbedding with dynamic wires is applied
+        after other non-AmplitudeEmbedding ops."""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def f(w):
+            qml.H(0)
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), [w])
+
+        with pytest.raises(
+            TransformError, match="Cannot apply qml.AmplitudeEmbedding with dynamic wires"
+        ):
+            _ = jax.make_jaxpr(f)(0)
+
 
 def test_plxpr_to_plxpr_transform():
     """Test that the plxpr transform works correctly for a simple example."""
 
     def qfunc():
-        qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+        qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
         qml.Hadamard(wires=0)
-        qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+        qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
         return qml.expval(qml.Z(0))
 
     jaxpr = jax.make_jaxpr(qfunc)()
@@ -463,9 +550,9 @@ def test_measure_prim(self):
 
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.Hadamard(wires=2)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
             qml.measure(wires=0)
 
         jaxpr = jax.make_jaxpr(qfunc)()
@@ -484,9 +571,9 @@ def test_measure_prim_block(self):
         @MergeAmplitudeEmbeddingInterpreter()
         def qfunc():
             qml.Hadamard(wires=2)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.measure(wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
         jaxpr = jax.make_jaxpr(qfunc)()
         collector = CollectOpsandMeas()
@@ -499,17 +586,31 @@ def qfunc():
         assert jaxpr.eqns[1].primitive == qml.Hadamard._primitive
         assert jaxpr.eqns[2].primitive == measure_prim
 
+    def test_dynamic_wires_embedding_after_measure(self):
+        """Test that an error is raised if an AmplitudeEmbedding with dynamic wires is applied after a
+        qml.measure."""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def qfunc(w):
+            qml.measure(wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=w)
+
+        with pytest.raises(
+            TransformError, match="Cannot apply qml.AmplitudeEmbedding with dynamic wires"
+        ):
+            _ = jax.make_jaxpr(qfunc)(1)
+
     def test_ctrl_transform_prim(self):
         """Test that the transform works correctly when applied with ctrl_transform_prim."""
 
         def ctrl_fn():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
             qml.X(0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=2)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=2)
 
         @MergeAmplitudeEmbeddingInterpreter()
         def f():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.ctrl(ctrl_fn, [3, 4])()
             qml.RY(0, 1)
 
@@ -532,9 +633,9 @@ def test_adjoint_transform_prim(self, lazy):
         @MergeAmplitudeEmbeddingInterpreter()
         def f():
             def g():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
                 qml.X(0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
             qml.adjoint(g, lazy=lazy)()
 
@@ -556,8 +657,8 @@ def test_cond_prim_only_true_branch(self):
         def f(x):
             @qml.cond(x > 2)
             def cond_f():
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
             cond_f()
 
@@ -579,22 +680,22 @@ def f(x):
             @qml.cond(x > 2)
             def cond_f():
                 qml.Z(0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=2)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=2)
                 return qml.expval(qml.Z(0))
 
             @cond_f.else_if(x > 1)
             def _():
                 qml.Y(1)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=2)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=2)
                 return qml.expval(qml.Y(0))
 
             @cond_f.otherwise
             def _():
                 qml.X(2)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
                 return qml.expval(qml.X(0))
 
             out = cond_f()
@@ -656,9 +757,9 @@ def f(n):
             # pylint:disable=unused-argument
             @qml.for_loop(n)
             def h(i):
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
                 qml.Hadamard(0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
 
             h()
 
@@ -680,9 +781,9 @@ def f(n):
 
             @qml.while_loop(lambda i: i < n)
             def h(i):
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
                 qml.Hadamard(0)
-                qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+                qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
                 return i + 1
 
             h(0)
@@ -704,9 +805,9 @@ def test_qnode_prim(self):
         @MergeAmplitudeEmbeddingInterpreter()
         @qml.qnode(dev)
         def circuit():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.Hadamard(0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
             return qml.expval(qml.Z(0))
 
         jaxpr = jax.make_jaxpr(circuit)()
@@ -728,9 +829,9 @@ def test_grad_prim(self):
         @MergeAmplitudeEmbeddingInterpreter()
         @qml.qnode(dev)
         def circuit(a, b):
-            qml.AmplitudeEmbedding(jax.numpy.array([a, b]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([a, b]), wires=0)
             qml.Hadamard(0)
-            qml.AmplitudeEmbedding(jax.numpy.array([a, b]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([a, b]), wires=1)
             return qml.expval(qml.Z(0))
 
         f = qml.grad(circuit)
@@ -760,9 +861,9 @@ def test_jacobian_prim(self):
         @MergeAmplitudeEmbeddingInterpreter()
         @qml.qnode(dev)
         def circuit(a, b):
-            qml.AmplitudeEmbedding(jax.numpy.array([a, b]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([a, b]), wires=0)
             qml.Hadamard(0)
-            qml.AmplitudeEmbedding(jax.numpy.array([a, b]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([a, b]), wires=1)
             return qml.expval(qml.Z(0))
 
         f = qml.jacobian(circuit)
@@ -788,6 +889,181 @@ def circuit(a, b):
         assert qfunc_jaxpr.eqns[2].primitive == qml.PauliZ._primitive
 
 
+class TestDynamicWiresControlFlowPrimitivesIntegration:
+    """Unit tests for using dynamic wires with control flow primitives."""
+
+    @pytest.mark.parametrize("embedding_location", ["cond", "after"])
+    @pytest.mark.parametrize("dynamic_wires", ["op", "embedding", "both"])
+    def test_cond_op_before(self, embedding_location, dynamic_wires):
+        """Test that applying an AmplitudeEmbedding with dynamic wires inside or after a cond after other
+        ops are already applied raises an error."""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(x, w):
+            op_wires = w if dynamic_wires in ("op", "both") else 0
+            embedding_wires = w if dynamic_wires in ("embedding", "both") else 0
+            qml.H(op_wires)
+
+            @qml.cond(x < 2)
+            def cond_fn():
+                if embedding_location == "cond":
+                    qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=embedding_wires)
+                qml.H(0)
+
+            @cond_fn.otherwise
+            def _():
+                qml.Y(0)
+
+            cond_fn()
+            if embedding_location == "after":
+                qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=embedding_wires)
+            return qml.expval(qml.Z(0))
+
+        with pytest.raises(TransformError, match="it is indeterminable if the wires overlap"):
+            _ = jax.make_jaxpr(circuit)(1.5, 0)
+
+    @pytest.mark.parametrize("dynamic_wires", ["op", "embedding", "both"])
+    def test_cond_dyn_wires_one_branch_embedding_other_branch(self, dynamic_wires):
+        """Test that an operator in one branch, and an AmplitudeEmbedding in another, where either has
+        dynamic wires, does not raise an error."""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(x, w):
+            op_wires = w if dynamic_wires in ("op", "both") else 0
+            embedding_wires = w if dynamic_wires in ("embedding", "both") else 0
+
+            @qml.cond(x < 2)
+            def cond_fn():
+                qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=embedding_wires)
+
+            @cond_fn.otherwise
+            def _():
+                qml.Y(op_wires)
+
+            cond_fn()
+            return qml.expval(qml.Z(0))
+
+        # No error should be raised
+        _ = jax.make_jaxpr(circuit)(1.5, 0)
+
+    @pytest.mark.parametrize("dynamic_wires", ["op", "embedding", "both"])
+    def test_cond_op_inside_embedding_after(self, dynamic_wires):
+        """Test that applying an AmplitudeEmbedding after a cond that contains an op where
+        either can have dynamic wires raises an error"""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(x, w):
+            op_wires = w if dynamic_wires in ("op", "both") else 0
+            embedding_wires = w if dynamic_wires in ("embedding", "both") else 0
+
+            @qml.cond(x < 2)
+            def cond_fn():
+                qml.H(op_wires)
+
+            cond_fn()
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), embedding_wires)
+            return qml.expval(qml.Z(0))
+
+        with pytest.raises(TransformError, match="it is indeterminable if the wires overlap"):
+            _ = jax.make_jaxpr(circuit)(1.5, 0)
+
+    @pytest.mark.parametrize("embedding_location", ["loop", "after"])
+    @pytest.mark.parametrize("dynamic_wires", ["op", "embedding", "both"])
+    def test_for_loop_op_before(self, dynamic_wires, embedding_location):
+        """Test that an error is raised if an op is applied before a for loop with an embedding inside or after
+        the loop where either can have dynamic wires"""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(n, w):
+            op_wires = w if dynamic_wires in ("op", "both") else 0
+            embedding_wires = w if dynamic_wires in ("embedding", "both") else 0
+            qml.H(op_wires)
+
+            @qml.for_loop(n)
+            def loop_fn(i):  # pylint: disable=unused-argument
+                if embedding_location == "loop":
+                    qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=embedding_wires)
+                qml.H(0)
+
+            loop_fn()
+            if embedding_location == "after":
+                qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=embedding_wires)
+            return qml.expval(qml.Z(0))
+
+        with pytest.raises(TransformError, match="it is indeterminable if the wires overlap"):
+            _ = jax.make_jaxpr(circuit)(5, 0)
+
+    @pytest.mark.parametrize("dynamic_wires", ["op", "embedding", "both"])
+    def test_for_loop_op_inside_embedding_after(self, dynamic_wires):
+        """Test that an error is raised if an op is applied inside a for loop with an embedding after
+        the loop where either can have dynamic wires"""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(n, w):
+            op_wires = w if dynamic_wires in ("op", "both") else 0
+            embedding_wires = w if dynamic_wires in ("embedding", "both") else 0
+
+            @qml.for_loop(n)
+            def loop_fn(i):  # pylint: disable=unused-argument
+                qml.H(op_wires)
+
+            loop_fn()
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), embedding_wires)
+            return qml.expval(qml.Z(0))
+
+        with pytest.raises(TransformError, match="it is indeterminable if the wires overlap"):
+            _ = jax.make_jaxpr(circuit)(5, 0)
+
+    @pytest.mark.parametrize("embedding_location", ["loop", "after"])
+    @pytest.mark.parametrize("dynamic_wires", ["op", "embedding", "both"])
+    def test_while_loop_op_before(self, dynamic_wires, embedding_location):
+        """Test that an error is raised if an op is applied before a while loop with an embedding inside or after
+        the loop where either can have dynamic wires"""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(x, w):
+            op_wires = w if dynamic_wires in ("op", "both") else 0
+            embedding_wires = w if dynamic_wires in ("embedding", "both") else 0
+            qml.H(op_wires)
+
+            @qml.while_loop(lambda arg: arg < 2)
+            def loop_fn(arg):
+                if embedding_location == "loop":
+                    qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=embedding_wires)
+                qml.H(0)
+                return arg - 1
+
+            loop_fn(x)
+            if embedding_location == "after":
+                qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), wires=embedding_wires)
+            return qml.expval(qml.Z(0))
+
+        with pytest.raises(TransformError, match="it is indeterminable if the wires overlap"):
+            _ = jax.make_jaxpr(circuit)(1.5, 0)
+
+    @pytest.mark.parametrize("dynamic_wires", ["op", "embedding", "both"])
+    def test_while_loop_op_inside_embedding_after(self, dynamic_wires):
+        """Test that an error is raised if an op is applied inside a while loop with an embedding after
+        the loop where either can have dynamic wires"""
+
+        @MergeAmplitudeEmbeddingInterpreter()
+        def circuit(x, w):
+            op_wires = w if dynamic_wires in ("op", "both") else 0
+            embedding_wires = w if dynamic_wires in ("embedding", "both") else 0
+
+            @qml.while_loop(lambda arg: arg < 2)
+            def loop_fn(arg):
+                qml.H(op_wires)
+                return arg - 1
+
+            loop_fn(x)
+            qml.AmplitudeEmbedding(jnp.array([1.0, 0.0]), embedding_wires)
+            return qml.expval(qml.Z(0))
+
+        with pytest.raises(TransformError, match="it is indeterminable if the wires overlap"):
+            _ = jax.make_jaxpr(circuit)(1.5, 0)
+
+
 class TestExpandPlxprTransformIntegration:
     """Test that the transform works with expand_plxpr_transform"""
 
@@ -796,9 +1072,9 @@ def test_example(self):
 
         @qml.transforms.optimization.merge_amplitude_embedding
         def qfunc():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.Hadamard(wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
             return qml.expval(qml.Z(0))
 
         jaxpr = jax.make_jaxpr(qfunc)()
@@ -823,9 +1099,9 @@ def test_decorator(self):
         @qml.capture.expand_plxpr_transforms
         @qml.transforms.optimization.merge_amplitude_embedding
         def qfunc():
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=0)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=0)
             qml.Hadamard(wires=0)
-            qml.AmplitudeEmbedding(jax.numpy.array([0.0, 1.0]), wires=1)
+            qml.AmplitudeEmbedding(jnp.array([0.0, 1.0]), wires=1)
             return qml.expval(qml.Z(0))
 
         jaxpr = jax.make_jaxpr(qfunc)()
diff --git a/tests/capture/transforms/test_capture_merge_rotations.py b/tests/capture/transforms/test_capture_merge_rotations.py
index c6c08102cc6..b867151c363 100644
--- a/tests/capture/transforms/test_capture_merge_rotations.py
+++ b/tests/capture/transforms/test_capture_merge_rotations.py
@@ -51,8 +51,8 @@ def test_traced_arguments(self):
         @MergeRotationsInterpreter()
         def f(a, b, c, wires):
             qml.RX(a, wires=wires)
-            qml.RY(a, wires=2)
             qml.RX(b, wires=wires)
+            qml.RY(a, wires=2)
             qml.Rot(0, 0, c, wires=1)
             qml.Rot(0, 0, c, wires=1)
             return qml.expval(qml.PauliZ(0))
@@ -413,6 +413,100 @@ def f():
         ops = collector.state["ops"]
         assert ops == expected_ops
 
+    def test_dynamic_wires_between_static_wires(self):
+        """Test that operations with dynamic wires between operations with static
+        wires cause merging to not happen."""
+
+        @MergeRotationsInterpreter()
+        def f(x, y, w):
+            qml.RX(x, 0)
+            qml.RY(y, w)
+            qml.RX(x, 0)
+            qml.RY(y, w)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(2.5, 3.5, 0)
+
+        dyn_wire = 0
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+        expected_meas = [qml.expval(qml.Z(0))]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 0), qml.RX(2.5, 0), qml.RY(3.5, 0)]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+        dyn_wire = 1
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 1), qml.RX(2.5, 0), qml.RY(3.5, 1)]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+    def test_same_dyn_wires_merge(self):
+        """Test that ops on the same dynamic wires get merged."""
+
+        @MergeRotationsInterpreter()
+        def f(x, y, w):
+            qml.RX(x, 0)
+            qml.RY(y, w)
+            qml.RY(y, w)
+            qml.RX(x, 0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(2.5, 3.5, 0)
+
+        dyn_wire = 0
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(jnp.array(7.0), 0), qml.RX(2.5, 0)]
+        expected_meas = [qml.expval(qml.Z(0))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+    def test_different_dyn_wires_interleaved(self):
+        """Test that ops on different dynamic wires interleaved with each other
+        do not merge."""
+
+        @MergeRotationsInterpreter()
+        def f(x, y, w1, w2):
+            qml.RX(x, w1)
+            qml.RY(y, w2)
+            qml.RX(x, w1)
+            qml.RY(y, w2)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(2.5, 3.5, 0, 0)
+
+        dyn_wires = (0, 0)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, *dyn_wires)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 0), qml.RX(2.5, 0), qml.RY(3.5, 0)]
+        expected_meas = [qml.expval(qml.Z(0))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+        dyn_wires = (0, 1)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, *dyn_wires)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 1), qml.RX(2.5, 0), qml.RY(3.5, 1)]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
 
 @pytest.mark.parametrize(("theta1, theta2"), [(0.1, 0.2), (0.1, -0.1)])
 def test_merge_rotations_plxpr_to_plxpr_transform(theta1, theta2):
diff --git a/tests/capture/transforms/test_capture_single_qubit_fusion.py b/tests/capture/transforms/test_capture_single_qubit_fusion.py
index 4cdb57e61f5..9c2305d3268 100644
--- a/tests/capture/transforms/test_capture_single_qubit_fusion.py
+++ b/tests/capture/transforms/test_capture_single_qubit_fusion.py
@@ -31,6 +31,7 @@
     while_loop_prim,
 )
 from pennylane.tape.plxpr_conversion import CollectOpsandMeas
+from pennylane.transforms.optimization.optimization_utils import fuse_rot_angles
 from pennylane.transforms.optimization.single_qubit_fusion import (
     SingleQubitFusionInterpreter,
     single_qubit_plxpr_to_plxpr,
@@ -340,12 +341,12 @@ def circuit(w1, w2, w3, w4):
         assert len(jaxpr_ops) == 6
 
         expected_ops = [
-            qml.Rot(np.pi / 2, np.pi, -np.pi / 2, wires=[0]),  # X not fused converted to Rot
+            qml.X(wires=[0]),  # X not fused
             qml.Rot(0, np.pi, np.pi, wires=[1]),  # Y and Z fused
             qml.CNOT(wires=[0, 1]),
-            qml.Rot(np.pi, 0, 0, wires=[2]),  # Z not fused converted to Rot
-            qml.CNOT(wires=[0, 2]),
+            qml.Z(wires=[2]),  # Z not fused
             qml.Hadamard(wires=[3]),  # H not fused and not converted to Rot
+            qml.CNOT(wires=[0, 2]),
         ]
 
         for op1, op2 in zip(jaxpr_ops, expected_ops, strict=True):
@@ -390,6 +391,131 @@ def circuit(params):
         for op1, op2 in zip(jaxpr_ops, expected_ops, strict=True):
             qml.assert_equal(op1, op2, check_interface=False)
 
+    def test_dynamic_wires_between_static_wires(self):
+        """Test that operations with dynamic wires between operations with static
+        wires cause fusion to not happen."""
+
+        @SingleQubitFusionInterpreter()
+        def f(x, y, w):
+            qml.RX(x, 0)
+            qml.RY(y, w)
+            qml.RX(x, 0)
+            qml.RY(y, w)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(2.5, 3.5, 0)
+
+        dyn_wire = 0
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+        expected_meas = [qml.expval(qml.Z(0))]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 0), qml.RX(2.5, 0), qml.RY(3.5, 0)]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+        dyn_wire = 1
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 1), qml.RX(2.5, 0), qml.RY(3.5, 1)]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+    def test_same_dyn_wires_fuse(self):
+        """Test that ops on the same dynamic wires get fused."""
+
+        @SingleQubitFusionInterpreter()
+        def f(x, y, w):
+            qml.RX(x, 0)
+            qml.RY(y, w)
+            qml.RY(y, w)
+            qml.RX(x, 0)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(2.5, 3.5, 0)
+
+        dyn_wire = 0
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, dyn_wire)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        # We use jax.jit here to make `phi` abstract. Otherwise, a different pipeline
+        # is used to compute the fused angles, which leads to the true and expected
+        # angles to be different.
+        @jax.jit
+        def get_double_ry_rot_angles(angles):
+            return fuse_rot_angles(angles, angles)
+
+        angles = jax.numpy.array(qml.RY(3.5, 0).single_qubit_rot_angles())
+        fused_angles = get_double_ry_rot_angles(angles)
+        expected_ops = [qml.RX(2.5, 0), qml.Rot(*fused_angles, 0), qml.RX(2.5, 0)]
+        expected_meas = [qml.expval(qml.Z(0))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+    def test_different_dyn_wires_interleaved(self):
+        """Test that ops on different dynamic wires interleaved with each other
+        do not fuse."""
+
+        @SingleQubitFusionInterpreter()
+        def f(x, y, w1, w2):
+            qml.RX(x, w1)
+            qml.RY(y, w2)
+            qml.RX(x, w1)
+            qml.RY(y, w2)
+            return qml.expval(qml.Z(0))
+
+        jaxpr = jax.make_jaxpr(f)(2.5, 3.5, 0, 0)
+
+        dyn_wires = (0, 0)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, *dyn_wires)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 0), qml.RX(2.5, 0), qml.RY(3.5, 0)]
+        expected_meas = [qml.expval(qml.Z(0))]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+        dyn_wires = (0, 1)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, 2.5, 3.5, *dyn_wires)
+        ops = collector.state["ops"]
+        meas = collector.state["measurements"]
+
+        expected_ops = [qml.RX(2.5, 0), qml.RY(3.5, 1), qml.RX(2.5, 0), qml.RY(3.5, 1)]
+        assert ops == expected_ops
+        assert meas == expected_meas
+
+    def test_diff_dyn_wires_non_fusible_op(self):
+        """Test that a non-fusible op with dynamic wires will cause all previous ops to be applied"""
+
+        @SingleQubitFusionInterpreter()
+        def f(w1, w2):
+            qml.H(w1)
+            qml.Z(0)
+            qml.CNOT([0, w2])
+
+        dyn_wires = (0, 1)
+        jaxpr = jax.make_jaxpr(f)(*dyn_wires)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, *dyn_wires)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.H(0),
+            qml.Rot(np.pi, 0, 0, 0),  # Z not fused but changed to Rot
+            qml.CNOT([0, 1]),
+        ]
+        assert ops == expected_ops
+
 
 class TestSingleQubitFusionHigherOrderPrimitives:
     """Unit tests for the single-qubit fusion transformation with higher order primitives."""


Reviews:
user: login: andrijapau
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T14:04:28Z
user: login: PietropaoloFrisoni
body: Thanks for the work on this Mudit! 

I have a few questions. Maybe it is best to start with a concrete example that right now doesn't provide the correct result and start discussing from there : )
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T17:20:10Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T21:44:04Z
user: login: PietropaoloFrisoni
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T14:12:28Z
user: login: PietropaoloFrisoni
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T14:17:05Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T19:17:11Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T19:18:24Z
user: login: PietropaoloFrisoni
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T19:42:30Z
user: login: andrijapau
body: Looks good to me. Probably should have adopted this approach when we were first implementing this transforms 🤣 . Thanks for the quick fix @mudit2812.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-26T14:03:26Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T17:56:52Z
user: login: isaacdevlugt
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-30T15:43:21Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-30T17:24:49Z
user: login: isaacdevlugt
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-30T17:41:11Z
user: login: andrijapau
body: Awesome, nice job with the tests! The parameterization is a bit weird to follow at first but it makes sense 🤣  

Can't believe how much more robust the transforms are now.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-30T17:44:17Z
user: login: PietropaoloFrisoni
body: Thanks Mudit, I think this is an amazing work :heart_eyes: 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-30T17:54:20Z
user: login: mudit2812
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-30T18:40:01Z
user: login: andrijapau
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-30T18:49:16Z

Review comments:
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
+
+    Transforming a circuit with dynamic wires using tape transorms may not raise an error, potentially
path: doc/news/program_capture_sharp_bits.rst
user: login: andrijapau
body: ```suggestion
    Transforming a circuit with dynamic wires using tape transforms may not raise an error, potentially
```
created_at: 2025-05-21T14:04:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: PietropaoloFrisoni
body: Can you please show an example that would lead to an incorrect result?
created_at: 2025-05-21T17:18:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: mudit2812
body: I can't actually think of an example of a transform that is readily available in pennylane that would cause this issue. My go-to's would've been `cancel_inverses` and `merge_rotations`, but those already have plxpr implementations 😅 
created_at: 2025-05-22T21:44:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -89,13 +89,20 @@ def _handle_non_fusible_op(self, op: Operator) -> list:
             """Handle an operation that cannot be fused into a Rot gate."""
 
             previous_ops_on_wires = self._retrieve_prev_ops_same_wire(op)
+            dyn_wires = {w for w in op.wires if qml.math.is_abstract(w)}
+            other_saved_wires = set(self.previous_ops.keys()) - dyn_wires
+
+            if dyn_wires and other_saved_wires:
+                # We cannot guarantee that ops on other static or dynamic wires won't have wire
+                # overlap with the current op, so we need to interpret all of them.
+                self.interpret_all_previous_ops()
path: pennylane/transforms/optimization/single_qubit_fusion.py
user: login: PietropaoloFrisoni
body: Yeah, this seems the only possible solution at the moment since we need to preserve correctness above optimization. This raises the question of how useful these (optimization) transforms are in the presence of dynamic wires (possibly combined with static wires) and how much we care about this. For the present context and the need to fix this bug, I totally agree with the strategy you adopted
created_at: 2025-05-23T14:12:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: PietropaoloFrisoni
body: Thanks for the clarification! So if I understand correctly, we anticipate that all transforms with plxpr implementation (like `cancel_inverses, merge_rotations`, etc.) provide a correct result even with dynamic wires at the cost of not optimizing the circuit. Is that so? If yes, since these transforms are not tape-based, why are we mentioning this here in the sharp bits page?
created_at: 2025-05-23T14:17:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: mudit2812
body: > we anticipate that all transforms with plxpr implementation (like cancel_inverses, merge_rotations, etc.) provide a correct result even with dynamic wires at the cost of not optimizing the circuit. Is that so?

Yes, this is correct.

> If yes, since these transforms are not tape-based, why are we mentioning this here in the sharp bits page?

This note in the sharp bits page is for tape-based transforms. What I meant was that I cannot think of a readily available tape-based transform that would satisfy the other constraints (no dynamic control flow, no post-processing, etc.) and will give incorrect results when dynamic wires are present.
created_at: 2025-05-23T19:17:11Z
author_association: CONTRIBUTOR
diff_hunk: @@ -89,13 +89,20 @@ def _handle_non_fusible_op(self, op: Operator) -> list:
             """Handle an operation that cannot be fused into a Rot gate."""
 
             previous_ops_on_wires = self._retrieve_prev_ops_same_wire(op)
+            dyn_wires = {w for w in op.wires if qml.math.is_abstract(w)}
+            other_saved_wires = set(self.previous_ops.keys()) - dyn_wires
+
+            if dyn_wires and other_saved_wires:
+                # We cannot guarantee that ops on other static or dynamic wires won't have wire
+                # overlap with the current op, so we need to interpret all of them.
+                self.interpret_all_previous_ops()
path: pennylane/transforms/optimization/single_qubit_fusion.py
user: login: mudit2812
body: I don't think we have any other solution. The same loss of optimization also already occurs in catalyst with dynamic wires, so I don't really see this as an incredibly significant loss.
created_at: 2025-05-23T19:18:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: PietropaoloFrisoni
body: Got it, my confusion came from the fact that I forgot that tape-based transforms are theoretically supported with program capture (although they are not 'natively' supported). 

I would probably use the bullet point or the warning in the sharp-bits page, since having both of them seems a little too vocal. What do you think?
created_at: 2025-05-23T19:42:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -75,11 +77,23 @@ def interpret_operation(self, op: Operator) -> None:
             so the output will not affect later equations in the circuit.
 
             """
+
             if not isinstance(op, AmplitudeEmbedding):
+                if any(is_abstract(w) for w in op.wires):
+                    if self.input_wires:
+                        self._merge_and_insert_at_the_start()
+                    self.interpret_all_previous_ops()
+                    self.dynamic_wires_encountered = True
+
                 self.previous_ops.append(op)
                 self.state["visited_wires"] = self.state["visited_wires"].union(set(op.wires))
                 return
 
+            if self.dynamic_wires_encountered:
+                raise TransformError(
+                    "Cannot use qml.AmplitudeEmbedding after operators with dynamic wires."
path: pennylane/transforms/optimization/merge_amplitude_embedding.py
user: login: andrijapau
body: Could we add a brief reason why? Could save some frustration on the user's side. 🤔 
```suggestion
                    "Cannot apply qml.AmplitudeEmbedding after operators with dynamic wires as it is indeterminable if the wires collide."
```
created_at: 2025-05-26T13:39:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: mudit2812
body: I'll tag @isaacdevlugt for this. My motivation for having it as a bullet point and also as a warning is that if the other statements in the other bullet points are true, then there will be an error. With this new bullet point, error is not guaranteed. Instead, a user might get incorrect answers.
created_at: 2025-05-28T17:56:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
path: doc/news/program_capture_sharp_bits.rst
user: login: isaacdevlugt
body: ```suggestion
  in the circuit whose wires are dynamic parameters).
```
created_at: 2025-05-30T15:37:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: isaacdevlugt
body: @mudit2812 your motivation is 100% valid but I think the added bullet point suffices :) 
created_at: 2025-05-30T15:39:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -534,6 +534,10 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Documentation 📝</h3>
 
+* The entry in the :doc:`/news/program_capture_sharp_bits` page for transforms has been updated. Using
+  dynamic wires with non-native transforms can lead to incorrect transformations.
+  [(#7426)](https://github.com/PennyLaneAI/pennylane/pull/7426)
path: doc/releases/changelog-dev.md
user: login: isaacdevlugt
body: ```suggestion
* The entry in the :doc:`/news/program_capture_sharp_bits` page for transforms has been updated; non-native transforms being applied to QNodes wherein operators have
  dynamic wires can lead to incorrect results.
  [(#7426)](https://github.com/PennyLaneAI/pennylane/pull/7426)
```
created_at: 2025-05-30T15:41:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -549,6 +553,14 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Bug fixes 🐛</h3>
 
+* Fix a bug where certain program capture native transforms transformed circuits incorrectly when
+  dynamic wires were present in the circuit. The affected transforms were:
path: doc/releases/changelog-dev.md
user: login: isaacdevlugt
body: ```suggestion
* Fixed a bug where certain transforms with a native program capture implementation give incorrect results when
  dynamic wires were present in the circuit. The affected transforms were:
```
created_at: 2025-05-30T15:43:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -448,12 +448,19 @@ to be used with certain limitations:
 
 * Transforms that return multiple tapes are not supported.
 * Transforms that return non-trivial post-processing functions are not supported.
+* Tape transforms may give incorrect results if the circuit has dynamic wires (i.e. there are operators
+  in the circuit dependent on dynamic parameters).
 * Tape transforms will fail to execute if the transformed quantum function or QNode contains:
 
    * ``qml.cond`` with dynamic parameters as predicates.
    * ``qml.for_loop`` with dynamic parameters for ``start``, ``stop``, or ``step``.
    * ``qml.while_loop``.
 
+.. warning::
path: doc/news/program_capture_sharp_bits.rst
user: login: mudit2812
body: Removed :) 
created_at: 2025-05-30T17:24:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -16,9 +16,10 @@
 import pytest
 
 import pennylane as qml
-from pennylane.exceptions import DeviceError
+from pennylane.transforms.core import TransformError
 
 jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
path: tests/capture/transforms/test_capture_merge_amplitude_embedding.py
user: login: andrijapau
body: is it not possible to just do `jnp = jax.numpy`?
created_at: 2025-05-30T17:39:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -16,9 +16,10 @@
 import pytest
 
 import pennylane as qml
-from pennylane.exceptions import DeviceError
+from pennylane.transforms.core import TransformError
 
 jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
path: tests/capture/transforms/test_capture_merge_amplitude_embedding.py
user: login: mudit2812
body: It is, but it makes no difference either way, and I have no preference. I've seen both patterns be used in our tests.

For the sake of not having to run CI again, I'd rather just leave it as is 😆 
created_at: 2025-05-30T18:40:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -16,9 +16,10 @@
 import pytest
 
 import pennylane as qml
-from pennylane.exceptions import DeviceError
+from pennylane.transforms.core import TransformError
 
 jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
path: tests/capture/transforms/test_capture_merge_amplitude_embedding.py
user: login: andrijapau
body: Sure NP
created_at: 2025-05-30T18:49:16Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-05-13T21:18:27Z
body: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
author_association: CONTRIBUTOR
user: login: codecov[bot]
created_at: 2025-05-16T22:21:32Z
body: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7426?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
:white_check_mark: All modified and coverable lines are covered by tests.
:white_check_mark: Project coverage is 99.68%. Comparing base ([`417d531`](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/417d5311aa721f6efc9a1efee006bf079a92dfe2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)) to head ([`edf9b0b`](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/edf9b0b770dc3bbf331856a0a52112ab5f4bc135?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)).
:warning: Report is 403 commits behind head on master.

<details><summary>Additional details and impacted files</summary>



```diff
@@           Coverage Diff           @@
##           master    #7426   +/-   ##
=======================================
  Coverage   99.68%   99.68%           
=======================================
  Files         531      531           
  Lines       51528    51574   +46     
=======================================
+ Hits        51365    51411   +46     
  Misses        163      163           
```
</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/7426?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE