Repository: PennyLaneAI/pennylane (python)

Issue:
Insert only for certain gates when using qml.transforms.insert
### Feature details

The `qml.transforms.insert()` transform was recently added in https://github.com/PennyLaneAI/pennylane/pull/1795, allowing users to insert a specified single qubit gate into a circuit with three orientations:
- at the start of the circuit
- at the end of the circuit
- after every gate in the circuit

However, another useful option might be:
_insert a specified (single qubit) gate after every gate of a given type_

For example, if we know that a CNOT gate adds depolarizing noise on each of its qubits, then we might want to transform our circuit to add `DepolarizingChannel` _only_ after `CNOT`.

### Implementation

The implementation in `qml.transforms.insert` should be simple, requiring an `if`-statement checking the type of the operation in the circuit, e.g., around [here](https://github.com/PennyLaneAI/pennylane/blob/e7217f4f9384184b0612c621ffa791b81050945c/pennylane/transforms/insert_ops.py#L211).

The main question would be the best choice of UI. For example, perhaps the `position` argument of `qml.transforms.insert` could be either a string (for the cases above) or a list containing the operations we want to insert a gate on, e.g.,
```python
position = [qml.CNOT, qml.CZ]
```
Alternatively/as well, we could just restrict `position` to be a string or a PL operation class, e.g.,
```python
position = qml.CNOT
```

Also note that the above assumes we want to insert a specified gate _after_ gates in the circuit. We might want to be even more flexible and have the ability to insert a specified gate _before_ gates in the circuit. This may involve an additional keyword argument being added.

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

Discussed in comment: https://github.com/PennyLaneAI/pennylane/pull/1795#discussion_r736573839

Pull request description:
Add functionality to qml.transforms.insert
**Context:**
 As discussed in #1876, some quality of life enhancements are made.
**Description of the Change:**
`before` argument can be used to specify whether to insert operation before or after
`position` can now take operations to insert only after specific gates
**Benefits:**
more functionality
**Possible Drawbacks:**
none
**Related GitHub Issues:**
Closes #1876 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 31a9314cd..f0cef4ccb 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -60,6 +60,9 @@
 
 <h3>Improvements</h3>
 
+* Insert transform now supports adding operation after or before certain specific gates.
+  [(#1980)](https://github.com/PennyLaneAI/pennylane/pull/1980)
+
 * Interferometer is now a class with `shape` method.
   [(#1946)](https://github.com/PennyLaneAI/pennylane/pull/1946)
 
diff --git a/pennylane/transforms/insert_ops.py b/pennylane/transforms/insert_ops.py
index 68e294fa5..7ec72f172 100644
--- a/pennylane/transforms/insert_ops.py
+++ b/pennylane/transforms/insert_ops.py
@@ -24,13 +24,39 @@ from pennylane.tape import QuantumTape
 from pennylane.tape.tape import STATE_PREP_OPS
 from pennylane.transforms.qfunc_transforms import qfunc_transform
 
+# pylint: disable=too-many-branches
+
+
+def _check_position(position):
+    """Checks the position argument to determine if an operation or list of operations was provided."""
+    not_op = False
+    req_ops = False
+    if isinstance(position, list):
+        req_ops = position.copy()
+        for operation in req_ops:
+            try:
+                if operation.__base__ != Operation:
+                    not_op = True
+            except AttributeError:
+                not_op = True
+    elif not isinstance(position, list):
+        try:
+            if position.__base__ == Operation:
+                req_ops = [position]
+            else:
+                not_op = True
+        except AttributeError:
+            not_op = True
+    return not_op, req_ops
+
 
 @qfunc_transform
 def insert(
     circuit: Union[callable, QuantumTape, Device],
     op: Union[callable, Type[Operation]],
     op_args: Union[tuple, float],
-    position: str = "all",
+    position: Union[str, list, Type[Operation]] = "all",
+    before: bool = False,
 ) -> Union[callable, QuantumTape]:
     """Insert an operation into specified points in an input circuit.
 
@@ -50,10 +76,13 @@ def insert(
             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
             float
-        position (str): Specification of where to add the operation. Should be one of: ``"all"`` to
-            add the operation after all gates (except state preparations); ``"start"`` to add the
-            operation to all wires at the start of the circuit (but after state preparations);
-            ``"end"`` to add the operation to all wires at the end of the circuit.
+        position (str or PennyLane operation or list of operations): Specification of where to add the operation.
+            Should be one of: ``"all"`` to add the operation after all gates (except state preparations);
+            ``"start"`` to add the operation to all wires at the start of the circuit (but after state preparations);
+            ``"end"`` to add the operation to all wires at the end of the circuit;
+            list of operations to add the operation before or after depending on ``before``.
+        before (bool): Whether to add the operation before the given operation(s) in ``position``.
+            Default is ``False`` and the operation is inserted after.
 
     Returns:
         callable or QuantumTape or Device: the updated version of the input circuit or an updated
@@ -62,7 +91,7 @@ def insert(
     Raises:
         ValueError: if a single operation acting on multiple wires is passed to ``op``
         ValueError: if the requested ``position`` argument is not ``'start'``, ``'end'`` or
-            ``'all'``
+            ``'all'`` OR PennyLane Operation
 
     **Example:**
 
@@ -189,8 +218,12 @@ def insert(
     if not isinstance(op, FunctionType) and op.num_wires != 1:
         raise ValueError("Only single-qubit operations can be inserted into the circuit")
 
-    if position not in ("start", "end", "all"):
-        raise ValueError("Position must be either 'start', 'end', or 'all' (default)")
+    not_op, req_ops = _check_position(position)
+
+    if position not in ("start", "end", "all") and not_op:
+        raise ValueError(
+            "Position must be either 'start', 'end', or 'all' (default) OR a PennyLane operation or list of operations."
+        )
 
     if not isinstance(op_args, Sequence):
         op_args = [op_args]
@@ -205,11 +238,22 @@ def insert(
             op(*op_args, wires=w)
 
     for circuit_op in circuit.operations[num_preps:]:
-        apply(circuit_op)
+        if not before:
+            apply(circuit_op)
+
         if position == "all":
             for w in circuit_op.wires:
                 op(*op_args, wires=w)
 
+        if req_ops:
+            for operation in req_ops:
+                if operation == type(circuit_op):
+                    for w in circuit_op.wires:
+                        op(*op_args, wires=w)
+
+        if before:
+            apply(circuit_op)
+
     if position == "end":
         for w in circuit.wires:
             op(*op_args, wires=w)


Test patch:
diff --git a/tests/transforms/test_insert_ops.py b/tests/transforms/test_insert_ops.py
index 14337234e..ee18a67cb 100644
--- a/tests/transforms/test_insert_ops.py
+++ b/tests/transforms/test_insert_ops.py
@@ -51,10 +51,11 @@ class TestInsert:
         with pytest.raises(ValueError, match="Only single-qubit operations can be inserted into"):
             insert(qml.CNOT, [])(self.tape)
 
-    def test_invalid_position(self):
+    @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="ABC")(self.tape)
+            insert(qml.AmplitudeDamping, 0.4, position=pos)(self.tape)
 
     def test_start(self):
         """Test if the expected tape is returned when the start position is requested"""
@@ -110,6 +111,87 @@ class TestInsert:
         assert tape.observables[0].wires.tolist() == [0, 1]
         assert tape.measurements[0].return_type is Expectation
 
+    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)
+
+        with QuantumTape() as tape_exp:
+            qml.PhaseDamping(0.4, wires=0)
+            qml.RX(0.9, wires=0)
+            qml.PhaseDamping(0.4, wires=1)
+            qml.RY(0.4, wires=1)
+            qml.PhaseDamping(0.4, wires=0)
+            qml.PhaseDamping(0.4, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.PhaseDamping(0.4, wires=0)
+            qml.RY(0.5, wires=0)
+            qml.PhaseDamping(0.4, wires=1)
+            qml.RX(0.6, wires=1)
+            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        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) == 1
+        assert tape.observables[0].name == ["PauliZ", "PauliZ"]
+        assert tape.observables[0].wires.tolist() == [0, 1]
+        assert tape.measurements[0].return_type is Expectation
+
+    def test_operation_as_position(self):
+        """Test if expected tape is returned when an operation is passed in position"""
+        tape = insert(qml.PhaseDamping, 0.4, position=qml.RX, before=True)(self.tape)
+
+        with QuantumTape() as tape_exp:
+            qml.PhaseDamping(0.4, wires=0)
+            qml.RX(0.9, wires=0)
+            qml.RY(0.4, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(0.5, wires=0)
+            qml.PhaseDamping(0.4, wires=1)
+            qml.RX(0.6, wires=1)
+            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        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) == 1
+        assert tape.observables[0].name == ["PauliZ", "PauliZ"]
+        assert tape.observables[0].wires.tolist() == [0, 1]
+        assert tape.measurements[0].return_type is Expectation
+
+    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)
+
+        with QuantumTape() as tape_exp:
+            qml.RX(0.9, wires=0)
+            qml.PhaseDamping(0.4, wires=0)
+            qml.RY(0.4, wires=1)
+            qml.PhaseDamping(0.4, wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(0.5, wires=0)
+            qml.PhaseDamping(0.4, wires=0)
+            qml.RX(0.6, wires=1)
+            qml.PhaseDamping(0.4, wires=1)
+            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        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) == 1
+        assert tape.observables[0].name == ["PauliZ", "PauliZ"]
+        assert tape.observables[0].wires.tolist() == [0, 1]
+        assert tape.measurements[0].return_type is Expectation
+
     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)