Repository: PennyLaneAI/pennylane
Pull request #4667 (2023-10-12 13:36:52)

Transform on devices and update insert
**Description of the Change:**

- Transforms can now be added to the preprocessing of device. `dev = transform(dev)`
- Insert is updated to the new transform system

**Benefits:**

**Possible Drawbacks:**
Default mixed is not a new generation device, the main use cases of applying insert on a device is for default mixed (adding noise).
- Add a special patch in transform dispatcher in the case of default mixed (to be removed in the future)

**Related GitHub Issues:**


github-actions[bot]: 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.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4667?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`ea7e1d4`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/ea7e1d4e4c5039aec9d4b9c059f1829d7dd10630?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.63% compared to head [(`d8d3759`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4667?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.62%.

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


```diff
@@            Coverage Diff             @@
##           master    #4667      +/-   ##
==========================================
- Coverage   99.63%   99.62%   -0.01%     
==========================================
  Files         376      376              
  Lines       33672    33697      +25     
==========================================
+ Hits        33548    33572      +24     
- Misses        124      125       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4667?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/core/transform\_dispatcher.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4667?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY29yZS90cmFuc2Zvcm1fZGlzcGF0Y2hlci5weQ==) | `99.27% <96.00%> (-0.73%)` | :arrow_down: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4667?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).

rmoyard: [sc-43680]
rmoyard: @timmysilv I will update the transforms changelog during feature freeze to make sure all details are covered such as the `@partial`

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d500a55ed..068d357c9 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -60,6 +60,9 @@
   to the new transform program system.
   [(#4573)](https://github.com/PennyLaneAI/pennylane/pull/4573)
 
+* Transforms can be applied on devices following the new device API.
+ [(#4667)](https://github.com/PennyLaneAI/pennylane/pull/4667)
+
 * All quantum functions transforms are update to the new transform program system.
  [(#4439)](https://github.com/PennyLaneAI/pennylane/pull/4439)
 
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index 8e81d499b..6b7660a8a 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -58,7 +58,7 @@ class TransformDispatcher:
 
         self._qnode_transform = self.default_qnode_transform
 
-    def __call__(self, *targs, **tkwargs):
+    def __call__(self, *targs, **tkwargs):  # pylint: disable=too-many-return-statements
         obj = None
 
         if targs:
@@ -75,6 +75,11 @@ class TransformDispatcher:
 
         if isinstance(obj, qml.QNode):
             return self._qnode_transform(obj, targs, tkwargs)
+        # TODO: Remove with the previous device generation
+        if isinstance(obj, qml.Device):
+            return self._old_device_transform(obj, targs, tkwargs)
+        if isinstance(obj, qml.devices.Device):
+            return self._device_transform(obj, targs, tkwargs)
         if callable(obj):
             return self._qfunc_transform(obj, targs, tkwargs)
 
@@ -229,6 +234,61 @@ class TransformDispatcher:
 
         return qfunc_transformed
 
+    def _old_device_transform(self, original_device, targs, tkwargs):
+        """Apply the transform on a device"""
+        if self._expand_transform:
+            raise TransformError("Device transform does not support expand transforms.")
+        if self._is_informative:
+            raise TransformError("Device transform does not support informative transforms.")
+        if self._final_transform:
+            raise TransformError("Device transform does not support final transforms.")
+        new_dev = copy.deepcopy(original_device)
+        transform = self._transform
+
+        @new_dev.custom_expand
+        def new_expand_fn(self, tape, *args, **kwargs):  # pylint: disable=unused-variable
+            tapes, _ = transform(tape, *targs, **tkwargs)
+            tape = tapes[0]
+            return self.default_expand_fn(tape, *args, **kwargs)
+
+        return new_dev
+
+    def _device_transform(self, original_device, targs, tkwargs):
+        """Apply the transform on a device"""
+        if self._expand_transform:
+            raise TransformError("Device transform does not support expand transforms.")
+        if self._is_informative:
+            raise TransformError("Device transform does not support informative transforms.")
+        if self._final_transform:
+            raise TransformError("Device transform does not support final transforms.")
+
+        class TransformedDevice(type(original_device)):
+            """A transformed device with updated preprocess method."""
+
+            def __init__(self, original_device, transform):
+                for key, value in original_device.__dict__.items():
+                    self.__setattr__(key, value)
+                self.transform = transform
+                self._original_device = original_device
+
+            def __repr__(self):
+                return f"Transformed Device({original_device.__repr__()} with additional preprocess transform {self.transform})"
+
+            def preprocess(
+                self, config: qml.devices.ExecutionConfig = qml.devices.DefaultExecutionConfig
+            ):
+                """This function updates the original device transform program to be applied."""
+                program, config = self.original_device.preprocess(config)
+                program.push_back(TransformContainer(self.transform, targs, tkwargs))
+                return program, config
+
+            @property
+            def original_device(self):
+                """Return the original device."""
+                return self._original_device
+
+        return TransformedDevice(original_device, self._transform)
+
 
 class TransformContainer:
     """Class to store a quantum transform with its args, kwargs and classical co-transforms.  Use
diff --git a/pennylane/transforms/insert_ops.py b/pennylane/transforms/insert_ops.py
index 17fbd63ac..0e8ea6a03 100644
--- a/pennylane/transforms/insert_ops.py
+++ b/pennylane/transforms/insert_ops.py
@@ -14,15 +14,13 @@
 """
 Provides transforms for inserting operations into quantum circuits.
 """
-from collections.abc import Sequence
 from types import FunctionType
-from typing import Type, Union
+from typing import Type, Union, Callable, Sequence
 
 import pennylane as qml
-from pennylane import Device, apply
 from pennylane.operation import Operation
 from pennylane.tape import QuantumTape
-from pennylane.transforms.qfunc_transforms import qfunc_transform
+from pennylane.transforms.core import transform
 from pennylane.ops.op_math import Adjoint
 
 # pylint: disable=too-many-branches
@@ -51,14 +49,14 @@ def _check_position(position):
     return not_op, req_ops
 
 
-@qfunc_transform
+@transform
 def insert(
-    circuit: Union[callable, QuantumTape, Device],
+    tape: QuantumTape,
     op: Union[callable, Type[Operation]],
     op_args: Union[tuple, float],
     position: Union[str, list, Type[Operation]] = "all",
     before: bool = False,
-) -> Union[callable, QuantumTape]:
+) -> (Sequence[QuantumTape], Callable):
     """Insert an operation into specified points in an input circuit.
 
     Circuits passed through this transform will be updated to have the operation, specified by the
@@ -71,8 +69,7 @@ def insert(
     for more information).
 
     Args:
-        circuit (callable or QuantumTape or pennylane.Device): the input circuit to be transformed, or a
-            device
+        circuit (QuantumTape): the input circuit to be transformed.
         op (callable or Type[Operation]): the single-qubit operation, or sequence of operations
             acting on a single qubit, to be inserted into the circuit
         op_args (tuple or float): the arguments fed to the operation, either as a tuple or a single
@@ -214,7 +211,7 @@ def insert(
     # decompose templates and their adjoints (which fixes a bug in the tutorial_error_mitigation demo)
     # TODO: change this to be cleaner and more robust
     try:
-        circuit = circuit.expand(
+        tape = tape.expand(
             stop_at=lambda op: not hasattr(qml.templates, op.name) and not isinstance(op, Adjoint)
         )
     except qml.QuantumFunctionError as e:
@@ -235,34 +232,45 @@ def insert(
 
     if not isinstance(op_args, Sequence):
         op_args = [op_args]
-
-    for prep_op in circuit.operations[: circuit.num_preps]:
-        apply(prep_op)
+    new_operations = []
+    for prep_op in tape.operations[: tape.num_preps]:
+        new_operations.append(prep_op)
 
     if position == "start":
-        for w in circuit.wires:
-            op(*op_args, wires=w)
+        for w in tape.wires:
+            sub_tape = qml.tape.make_qscript(op)(*op_args, wires=w)
+            new_operations.extend(sub_tape.operations)
 
-    for circuit_op in circuit.operations[circuit.num_preps :]:
+    for circuit_op in tape.operations[tape.num_preps :]:
         if not before:
-            apply(circuit_op)
+            new_operations.append(circuit_op)
 
         if position == "all":
             for w in circuit_op.wires:
-                op(*op_args, wires=w)
+                sub_tape = qml.tape.make_qscript(op)(*op_args, wires=w)
+                new_operations.extend(sub_tape.operations)
 
         if req_ops:
             for operation in req_ops:
                 if operation == type(circuit_op):
                     for w in circuit_op.wires:
-                        op(*op_args, wires=w)
+                        sub_tape = qml.tape.make_qscript(op)(*op_args, wires=w)
+                        new_operations.extend(sub_tape.operations)
 
         if before:
-            apply(circuit_op)
+            new_operations.append(circuit_op)
 
     if position == "end":
-        for w in circuit.wires:
-            op(*op_args, wires=w)
+        for w in tape.wires:
+            sub_tape = qml.tape.make_qscript(op)(*op_args, wires=w)
+            new_operations.extend(sub_tape.operations)
+
+    new_tape = type(tape)(new_operations, tape.measurements, shots=tape.shots)
+
+    def null_postprocessing(results):
+        """A postprocesing function returned by a transform that only converts the batch of results
+        into a result for a single ``QuantumTape``.
+        """
+        return results[0]
 
-    for m in circuit.measurements:
-        apply(m)
+    return [new_tape], null_postprocessing
diff --git a/pennylane/transforms/qfunc_transforms.py b/pennylane/transforms/qfunc_transforms.py
index fd5b67145..0717e0215 100644
--- a/pennylane/transforms/qfunc_transforms.py
+++ b/pennylane/transforms/qfunc_transforms.py
@@ -151,7 +151,9 @@ class single_tape_transform:
         return qs
 
 
-def _create_qfunc_internal_wrapper(fn, tape_transform, transform_args, transform_kwargs):
+def _create_qfunc_internal_wrapper(
+    fn, tape_transform, transform_args, transform_kwargs
+):  # pragma: no cover
     """Convenience function to create the internal wrapper function
     generated by the qfunc_transform decorator"""
     if isinstance(fn, qml.Device):


Test patch:
diff --git a/tests/transforms/test_experimental/test_transform_dispatcher.py b/tests/transforms/test_experimental/test_transform_dispatcher.py
index bf9d18e6e..ec3a5afb3 100644
--- a/tests/transforms/test_experimental/test_transform_dispatcher.py
+++ b/tests/transforms/test_experimental/test_transform_dispatcher.py
@@ -530,3 +530,80 @@ class TestTransformDispatcher:  # pylint: disable=too-many-public-methods
         qnode = qml.QNode(transformed_qfunc, qml.device("default.qubit"))
         result = qnode()
         assert isinstance(result, type_)
+
+    @pytest.mark.parametrize("valid_transform", valid_transforms)
+    def test_device_transform(self, valid_transform):
+        """Test a device transform."""
+        dispatched_transform = transform(valid_transform)
+        new_dev = dispatched_transform(dev, index=0)
+
+        assert new_dev.original_device is dev
+        assert repr(new_dev).startswith("Transformed Device")
+
+        program, _ = dev.preprocess()
+        new_program, _ = new_dev.preprocess()
+
+        assert isinstance(program, qml.transforms.core.TransformProgram)
+        assert isinstance(new_program, qml.transforms.core.TransformProgram)
+
+        assert len(program) == 4
+        assert len(new_program) == 5
+
+        assert new_program[-1].transform is valid_transform
+
+    @pytest.mark.parametrize("valid_transform", valid_transforms)
+    def test_old_device_transform(self, valid_transform):
+        """Test a device transform on old device."""
+        dispatched_transform = transform(valid_transform)
+        device = qml.device("default.mixed", wires=2)
+        new_dev = dispatched_transform(device, index=0)
+
+        assert isinstance(new_dev, type(device))
+        assert new_dev.custom_expand_fn
+        assert repr(device).startswith("<DefaultMixed device (wires=2, shots=None)")
+        assert repr(new_dev).startswith("<DefaultMixed device (wires=2, shots=None)")
+
+    @pytest.mark.parametrize("valid_transform", valid_transforms)
+    def test_device_transform_error(self, valid_transform):
+        """Test that the device transform returns errors."""
+
+        with pytest.raises(
+            TransformError, match="Device transform does not support informative transforms."
+        ):
+            dispatched_transform = transform(valid_transform, is_informative=True)
+            dispatched_transform(dev, index=0)
+
+        with pytest.raises(
+            TransformError, match="Device transform does not support final transforms."
+        ):
+            dispatched_transform = transform(valid_transform, final_transform=True)
+            dispatched_transform(dev, index=0)
+
+        with pytest.raises(
+            TransformError, match="Device transform does not support expand transforms."
+        ):
+            dispatched_transform = transform(valid_transform, expand_transform=valid_transform)
+            dispatched_transform(dev, index=0)
+
+    @pytest.mark.parametrize("valid_transform", valid_transforms)
+    def test_old_device_transform_error(self, valid_transform):
+        """Test that the old device transform returns errors."""
+        device = qml.device("default.mixed", wires=2)
+
+        with pytest.raises(
+            TransformError, match="Device transform does not support informative transforms."
+        ):
+            dispatched_transform = transform(valid_transform, is_informative=True)
+            dispatched_transform(device, index=0)
+
+        with pytest.raises(
+            TransformError, match="Device transform does not support final transforms."
+        ):
+            dispatched_transform = transform(valid_transform, final_transform=True)
+            dispatched_transform(device, index=0)
+
+        with pytest.raises(
+            TransformError, match="Device transform does not support expand transforms."
+        ):
+            dispatched_transform = transform(valid_transform, expand_transform=valid_transform)
+            dispatched_transform(device, index=0)
diff --git a/tests/transforms/test_insert_ops.py b/tests/transforms/test_insert_ops.py
index 6b22ffda8..08742f5da 100644
--- a/tests/transforms/test_insert_ops.py
+++ b/tests/transforms/test_insert_ops.py
@@ -14,6 +14,7 @@
 """
 Tests for the insert transform.
 """
+from functools import partial
 import numpy as np
 import pytest
 
@@ -59,17 +60,18 @@ class TestInsert:
     def test_multiwire_op(self):
         """Tests if a ValueError is raised when multiqubit operations are requested"""
         with pytest.raises(ValueError, match="Only single-qubit operations can be inserted into"):
-            insert(qml.CNOT, [])(self.tape)
+            insert(self.tape, qml.CNOT, [])
 
     @pytest.mark.parametrize("pos", [1, ["all", qml.RY, int], "ABC", str])
     def test_invalid_position(self, pos):
         """Test if a ValueError is raised when an invalid position is requested"""
         with pytest.raises(ValueError, match="Position must be either 'start', 'end', or 'all'"):
-            insert(qml.AmplitudeDamping, 0.4, position=pos)(self.tape)
+            insert(self.tape, qml.AmplitudeDamping, 0.4, position=pos)
 
     def test_start(self):
         """Test if the expected tape is returned when the start position is requested"""
-        tape = insert(qml.AmplitudeDamping, 0.4, position="start")(self.tape)
+        tapes, _ = insert(self.tape, qml.AmplitudeDamping, 0.4, position="start")
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.AmplitudeDamping(0.4, wires=0)
@@ -95,7 +97,8 @@ class TestInsert:
 
     def test_all(self):
         """Test if the expected tape is returned when the all position is requested"""
-        tape = insert(qml.PhaseDamping, 0.4, position="all")(self.tape)
+        tapes, _ = insert(self.tape, qml.PhaseDamping, 0.4, position="all")
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.RX(0.9, wires=0)
@@ -125,8 +128,8 @@ class TestInsert:
 
     def test_before(self):
         """Test if the expected tape is returned when the before argument is True"""
-        tape = insert(qml.PhaseDamping, 0.4, position="all", before=True)(self.tape)
-
+        tapes, _ = insert(self.tape, qml.PhaseDamping, 0.4, position="all", before=True)
+        tape = tapes[0]
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.PhaseDamping(0.4, wires=0)
             qml.RX(0.9, wires=0)
@@ -158,7 +161,8 @@ class TestInsert:
     @pytest.mark.parametrize("op", op_lst)
     def test_operation_as_position(self, op):
         """Test if expected tape is returned when an operation is passed in position"""
-        tape = insert(qml.PhaseDamping, 0.4, position=op, before=True)(self.custom_tape)
+        tapes, _ = insert(self.custom_tape, qml.PhaseDamping, 0.4, position=op, before=True)
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             if op == qml.RX:
@@ -193,7 +197,8 @@ class TestInsert:
 
     def test_operation_list_as_position(self):
         """Test if expected tape is returned when an operation list is passed in position"""
-        tape = insert(qml.PhaseDamping, 0.4, position=[qml.RX, qml.RY])(self.tape)
+        tapes, _ = insert(self.tape, qml.PhaseDamping, 0.4, position=[qml.RX, qml.RY])
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.RX(0.9, wires=0)
@@ -221,7 +226,8 @@ class TestInsert:
 
     def test_end(self):
         """Test if the expected tape is returned when the end position is requested"""
-        tape = insert(qml.GeneralizedAmplitudeDamping, [0.4, 0.5], position="end")(self.tape)
+        tapes, _ = insert(self.tape, qml.GeneralizedAmplitudeDamping, [0.4, 0.5], position="end")
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.RX(0.9, wires=0)
@@ -248,7 +254,8 @@ class TestInsert:
     def test_start_with_state_prep(self):
         """Test if the expected tape is returned when the start position is requested in a tape
         that has state preparation"""
-        tape = insert(qml.AmplitudeDamping, 0.4, position="start")(self.tape_with_prep)
+        tapes, _ = insert(self.tape_with_prep, qml.AmplitudeDamping, 0.4, position="start")
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.StatePrep([1, 0], wires=0)
@@ -276,7 +283,8 @@ class TestInsert:
     def test_all_with_state_prep(self):
         """Test if the expected tape is returned when the all position is requested in a tape
         that has state preparation"""
-        tape = insert(qml.PhaseDamping, 0.4, position="all")(self.tape_with_prep)
+        tapes, _ = insert(self.tape_with_prep, qml.PhaseDamping, 0.4, position="all")
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.StatePrep([1, 0], wires=0)
@@ -308,9 +316,10 @@ class TestInsert:
     def test_end_with_state_prep(self):
         """Test if the expected tape is returned when the end position is requested in a tape
         that has state preparation"""
-        tape = insert(qml.GeneralizedAmplitudeDamping, [0.4, 0.5], position="end")(
-            self.tape_with_prep
+        tapes, _ = insert(
+            self.tape_with_prep, qml.GeneralizedAmplitudeDamping, [0.4, 0.5], position="end"
         )
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.StatePrep([1, 0], wires=0)
@@ -343,7 +352,8 @@ class TestInsert:
             qml.RX(x, wires=wires)
             qml.PhaseShift(y, wires=wires)
 
-        tape = insert(op, [0.4, 0.5], position="end")(self.tape)
+        tapes, _ = insert(self.tape, op=op, op_args=[0.4, 0.5], position="end")
+        tape = tapes[0]
 
         with qml.queuing.AnnotatedQueue() as q_tape_exp:
             qml.RX(0.9, wires=0)
@@ -375,8 +385,8 @@ def test_insert_qnode():
     without."""
     dev = qml.device("default.mixed", wires=2)
 
+    @partial(insert, op=qml.AmplitudeDamping, op_args=0.2, position="end")
     @qml.qnode(dev)
-    @insert(qml.AmplitudeDamping, 0.2, position="end")
     def f_noisy(w, x, y, z):
         qml.RX(w, wires=0)
         qml.RY(x, wires=1)
@@ -399,8 +409,66 @@ def test_insert_qnode():
     assert not np.isclose(f_noisy(*args), f(*args))
 
 
-def test_insert_dev(mocker):
-    """Test if a device transformed by the insert function does successfully add noise to
+def test_insert_dev():
+    """Test if an device transformed by the insert function does successfully add noise to
+    subsequent circuit executions"""
+    with qml.queuing.AnnotatedQueue() as q_in_tape:
+        qml.RX(0.9, wires=0)
+        qml.RY(0.4, wires=1)
+        qml.CNOT(wires=[0, 1])
+        qml.RY(0.5, wires=0)
+        qml.RX(0.6, wires=1)
+        qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+        qml.expval(qml.PauliZ(0))
+
+    in_tape = QuantumScript.from_queue(q_in_tape)
+    dev = qml.device("default.qubit", wires=2)
+    program, _ = dev.preprocess()
+    res_without_noise = qml.execute(
+        [in_tape], dev, qml.gradients.param_shift, transform_program=program
+    )
+
+    new_dev = insert(dev, qml.PhaseShift, 0.4)
+    new_program, _ = new_dev.preprocess()
+    tapes, _ = new_program([in_tape])
+    tape = tapes[0]
+    res_with_noise = qml.execute([in_tape], new_dev, qml.gradients, transform_program=new_program)
+
+    with qml.queuing.AnnotatedQueue() as q_tape_exp:
+        qml.RX(0.9, wires=0)
+        qml.PhaseShift(0.4, wires=0)
+        qml.RY(0.4, wires=1)
+        qml.PhaseShift(0.4, wires=1)
+        qml.CNOT(wires=[0, 1])
+        qml.PhaseShift(0.4, wires=0)
+        qml.PhaseShift(0.4, wires=1)
+        qml.RY(0.5, wires=0)
+        qml.PhaseShift(0.4, wires=0)
+        qml.RX(0.6, wires=1)
+        qml.PhaseShift(0.4, wires=1)
+        qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+        qml.expval(qml.PauliZ(0))
+
+    tape_exp = QuantumScript.from_queue(q_tape_exp)
+    assert all(o1.name == o2.name for o1, o2 in zip(tape.operations, tape_exp.operations))
+    assert all(o1.wires == o2.wires for o1, o2 in zip(tape.operations, tape_exp.operations))
+    assert all(
+        np.allclose(o1.parameters, o2.parameters)
+        for o1, o2 in zip(tape.operations, tape_exp.operations)
+    )
+    assert len(tape.measurements) == 2
+    assert tape.observables[0].name == ["PauliZ", "PauliZ"]
+    assert tape.observables[0].wires.tolist() == [0, 1]
+    assert tape.measurements[0].return_type is Expectation
+    assert tape.observables[1].name == "PauliZ"
+    assert tape.observables[1].wires.tolist() == [0]
+    assert tape.measurements[1].return_type is Expectation
+
+    assert not np.allclose(res_without_noise, res_with_noise)
+
+
+def test_insert_old_dev(mocker):
+    """Test if a old device transformed by the insert function does successfully add noise to
     subsequent circuit executions"""
     with qml.queuing.AnnotatedQueue() as q_in_tape:
         qml.RX(0.9, wires=0)
@@ -415,7 +483,7 @@ def test_insert_dev(mocker):
     dev = qml.device("default.mixed", wires=2)
     res_without_noise = qml.execute([in_tape], dev, qml.gradients.param_shift)
 
-    new_dev = insert(qml.PhaseDamping, 0.4)(dev)
+    new_dev = insert(dev, qml.PhaseDamping, 0.4)
     spy = mocker.spy(new_dev, "default_expand_fn")
 
     res_with_noise = qml.execute([in_tape], new_dev, qml.gradients.param_shift)
@@ -458,8 +526,8 @@ def test_insert_template():
     """Test that ops are inserted correctly into a decomposed template"""
     dev = qml.device("default.mixed", wires=2)
 
+    @partial(insert, op=qml.PhaseDamping, op_args=0.3, position="all")
     @qml.qnode(dev)
-    @insert(qml.PhaseDamping, 0.3, position="all")
     def f1(w1, w2):
         qml.SimplifiedTwoDesign(w1, w2, wires=[0, 1])
         return qml.expval(qml.PauliZ(0))
@@ -487,6 +555,7 @@ def test_insert_template():
 
 def test_insert_decorator_causes_custom_insert_error_non_qwc_obs():
     """Test that the insert transform catches and reports errors from the enclosed function."""
+
     # pylint: disable=unused-argument
 
     def noise(noise_param, wires):
@@ -496,7 +565,7 @@ def test_insert_decorator_causes_custom_insert_error_non_qwc_obs():
     dev = qml.device("default.mixed", wires=2)
 
     @qml.qnode(dev)
-    @qml.transforms.insert(noise, 0.3, position="all")
+    @partial(insert, op=noise, op_args=0.3, position="all")
     def noisy_circuit(circuit_param):
         qml.RY(circuit_param, wires=0)
         qml.Hadamard(wires=0)
diff --git a/tests/transforms/test_mitigate.py b/tests/transforms/test_mitigate.py
index 5a321fcc1..19fe155e1 100644
--- a/tests/transforms/test_mitigate.py
+++ b/tests/transforms/test_mitigate.py
@@ -114,7 +114,7 @@ class TestMitigateWithZNE:
         noise_strength = 0.05
 
         dev_noise_free = qml.device("default.mixed", wires=2)
-        dev = qml.transforms.insert(qml.AmplitudeDamping, noise_strength)(dev_noise_free)
+        dev = qml.transforms.insert(dev_noise_free, qml.AmplitudeDamping, noise_strength)
 
         n_wires = 2
         n_layers = 2
@@ -443,7 +443,7 @@ noise_gate = qml.PhaseDamping
 
 # Load devices
 dev_ideal = qml.device("default.mixed", wires=2)
-dev_noisy = qml.transforms.insert(noise_gate, 0.05)(dev_ideal)
+dev_noisy = qml.transforms.insert(dev_ideal, noise_gate, 0.05)
 
 out_ideal = np.sqrt(2) / 2 + np.sqrt(2)
 grad_ideal_0 = [-np.sqrt(2) / 2, -np.sqrt(2)]