transforms.insert: Update docs, improve decomposition (#5681)

**Context:**
The first example in the docs of `qml.transforms.insert` was reported to
not be working
[here](https://discuss.pennylane.ai/t/qml-device-error-with-qiskit-aer/4556/9).
Other examples also do not work.

**Description of the Change:**
This PR changes the docs examples to work.

It also replicates the change in #5424 to use
`devices.preprocess.decompose` instead of `tape.expand`, adding support
for non-commuting measurements (which is unrelated to the transform
anyways), and removing the need for a custom error to be raised.

**Benefits:**
Docs examples work.
Non-commuting measurements support.

**Possible Drawbacks:**
Similar to #5424, this introduces cross-dependency to the devices
module.

**Related GitHub Issues:**

doc/releases/changelog-dev.md

@@ -108,9 +108,10 @@
 * Finite shot circuits with a `qml.probs` measurement, both with a `wires` or `op` argument, can now be compiled with `jax.jit`.
   [(#5619)](https://github.com/PennyLaneAI/pennylane/pull/5619)
 
-* `param_shift`, `finite_diff`, `compile`, `merge_rotations`, and `transpile` now all work
-  with circuits with non-commuting measurements.
+* `param_shift`, `finite_diff`, `compile`, `insert`, `merge_rotations`, and `transpile` now
+  all work with circuits with non-commuting measurements.
   [(#5424)](https://github.com/PennyLaneAI/pennylane/pull/5424)
+  [(#5681)](https://github.com/PennyLaneAI/pennylane/pull/5681)
 
 * A correction is added to `bravyi_kitaev` to call the correct function for a FermiSentence input.
   [(#5671)](https://github.com/PennyLaneAI/pennylane/pull/5671)

pennylane/transforms/insert_ops.py

@@ -89,7 +89,6 @@ def insert(
 
 
     Raises:
-        QuantumFunctionError: if some observables in the tape are not qubit-wise commuting
         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'`` OR PennyLane Operation
@@ -100,9 +99,11 @@ def insert(
 
     .. code-block:: python3
 
+        from functools import partial
+
         dev = qml.device("default.mixed", wires=2)
 
-        @partial(qml.transforms.insert, qml.AmplitudeDamping, 0.2, position="end")
+        @partial(qml.transforms.insert, op=qml.AmplitudeDamping, op_args=0.2, position="end")
         @qml.qnode(dev)
         def f(w, x, y, z):
             qml.RX(w, wires=0)
@@ -141,7 +142,7 @@ def op(x, y, wires):
             dev = qml.device("default.qubit", wires=2)
 
             @qml.qnode(dev)
-            @qml.transforms.insert(op, [0.2, 0.3], position="end")
+            @partial(qml.transforms.insert, op=op, op_args=[0.2, 0.3], position="end")
             def f(w, x, y, z):
                 qml.RX(w, wires=0)
                 qml.RY(x, wires=1)
@@ -175,7 +176,7 @@ def f(w, x, y, z):
         We can add the :class:`~.AmplitudeDamping` channel to the end of the circuit using:
 
         >>> from pennylane.transforms import insert
-        >>> noisy_tape = insert(tape, qml.AmplitudeDamping, 0.05, position="end")
+        >>> [noisy_tape], _ = insert(tape, qml.AmplitudeDamping, 0.05, position="end")
         >>> print(qml.drawer.tape_text(noisy_tape, decimals=2))
         0: ──RX(0.90)─╭●──RY(0.50)──AmplitudeDamping(0.05)─┤ ╭<Z@Z>
         1: ──RY(0.40)─╰X──RX(0.60)──AmplitudeDamping(0.05)─┤ ╰<Z@Z>
@@ -210,17 +211,18 @@ def f(w, x, y, z):
         >>> qnode_noisy(0.9, 0.4, 0.5, 0.6)
         tensor(0.72945434, requires_grad=True)
     """
-    # 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:
-        tape = tape.expand(
-            stop_at=lambda op: not hasattr(qml.templates, op.name) and not isinstance(op, Adjoint)
-        )
-    except qml.QuantumFunctionError as e:
-        raise qml.QuantumFunctionError(
-            "The insert transform cannot transform a circuit measuring non-commuting observables. "
-            "Consider wrapping the gates in their own function and transforming only that function."
-        ) from e
+
+    # decompose templates and their adjoints to fix a bug in the tutorial_error_mitigation demo
+    def stop_at(obj):
+        if not isinstance(obj, qml.operation.Operator):
+            return True
+        if not obj.has_decomposition:
+            return True
+        return not (hasattr(qml.templates, obj.name) or isinstance(obj, Adjoint))
+
+    error_type = (qml.operation.DecompositionUndefinedError,)
+    decompose = qml.devices.preprocess.decompose
+    [tape], _ = decompose(tape, stopping_condition=stop_at, name="insert", error=error_type)
 
     if not isinstance(op, FunctionType) and op.num_wires != 1:
         raise ValueError("Only single-qubit operations can be inserted into the circuit")

tests/transforms/test_insert_ops.py

@@ -602,28 +602,32 @@ def f2(w1, w2):
     assert np.allclose(f1(w1, w2), f2(w1, w2))
 
 
-def test_insert_decorator_causes_custom_insert_error_non_qwc_obs():
+def test_insert_transform_works_with_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):
+    def op(noise_param, wires):
+        # pylint: disable=unused-argument
         qml.CRX(noise_param, wires=[0, 1])
         qml.CNOT(wires=[1, 0])
 
-    dev = qml.device("default.mixed", wires=2)
+    dev = qml.device("default.qubit", wires=2)
 
     @qml.qnode(dev)
-    @partial(insert, op=noise, op_args=0.3, position="all")
+    @partial(insert, op=op, op_args=0.3, position="all")
     def noisy_circuit(circuit_param):
         qml.RY(circuit_param, wires=0)
         qml.Hadamard(wires=0)
         qml.T(wires=0)
         return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0)), qml.expval(qml.PauliZ(0))
 
-    # This tape's expansion fails, but shouldn't cause a downstream IndexError. See issue #3103
-    with pytest.raises(
-        qml.QuantumFunctionError,
-        match="The insert transform cannot transform a circuit measuring non-commuting observables",
-    ):
-        noisy_circuit(0.4)
+    @qml.qnode(dev)
+    def explicit_circuit(circuit_param):
+        qml.RY(circuit_param, wires=0)
+        op(0.3, None)
+        qml.Hadamard(wires=0)
+        op(0.3, None)
+        qml.T(wires=0)
+        op(0.3, None)
+        return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliY(0)), qml.expval(qml.PauliZ(0))
+
+    assert np.allclose(noisy_circuit(0.4), explicit_circuit(0.4))