Repository: PennyLaneAI/pennylane
Pull request #4949 (2023-12-14 20:11:42)

Make `qml.simplify` a transform
**Context:**
`qml.simplify` can be applied to operations, measurements, tapes, qnodes, or qfuncs. It seems appropriate for it to be a transform. So I made it one.

**Description of the Change:**
* Added `_simplify_transform` with a `transform` decorator to handle transforming tapes, qnodes, and callables. The public `simplify` function directly handles operators and measurements, and dispatches to `_simplify_transform` for other valid inputs.

**Benefits:**
`qml.simplify` is consistent with other op functions/transforms.

**Possible Drawbacks:**
When given a tape, `qml.simplify` used to return a tape, but now it returns a list with one tape and a processing function. Small breaking change, but this was also a breaking change for the other op transforms, so I don't think this is a major deal.

**Related GitHub Issues:**

[sc-52122]

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/4949?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`5eb005e`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/5eb005ea247f67518bab1be127e930c1005597bd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.49% compared to head [(`2bf21c0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4949?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.49%.


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


```diff
@@            Coverage Diff             @@
##           master    #4949      +/-   ##
==========================================
- Coverage   99.49%   99.49%   -0.01%     
==========================================
  Files         388      388              
  Lines       35467    35196     -271     
==========================================
- Hits        35289    35017     -272     
- Misses        178      179       +1     
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index cb86dc86e..38fa43b9a 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -383,6 +383,9 @@
 
 <h4>Other improvements</h4>
 
+* `qml.simplify` now uses the new transforms API.
+  [(#4949)](https://github.com/PennyLaneAI/pennylane/pull/4949)
+
 * The formal requirement that type hinting be providing when using
   the `qml.transform` decorator has been removed. Type hinting can still
   be used, but is now optional. Please use a type checker such as
diff --git a/pennylane/ops/functions/simplify.py b/pennylane/ops/functions/simplify.py
index 7814c6151..51c5f30f0 100644
--- a/pennylane/ops/functions/simplify.py
+++ b/pennylane/ops/functions/simplify.py
@@ -15,15 +15,14 @@
 This module contains the qml.simplify function.
 """
 from copy import copy
-from functools import wraps
-from typing import Callable, Union
+from typing import Callable, Union, Sequence
 
 import pennylane as qml
 from pennylane.measurements import MeasurementProcess
 from pennylane.operation import Operator
 from pennylane.qnode import QNode
 from pennylane.queuing import QueuingManager
-from pennylane.tape import QuantumScript, make_qscript, QuantumTape
+from pennylane.tape import QuantumScript, QuantumTape
 
 
 def simplify(input: Union[Operator, MeasurementProcess, QuantumTape, QNode, Callable]):
@@ -31,11 +30,14 @@ def simplify(input: Union[Operator, MeasurementProcess, QuantumTape, QNode, Call
     or number of rotation parameters.
 
     Args:
-        input (.Operator, pennylane.QNode, .QuantumTape, or Callable): an operator, quantum node,
-            tape or function that applies quantum operations
+        input (.Operator, .MeasurementProcess, pennylane.QNode, .QuantumTape, or Callable): an
+            operator, quantum node, tape or function that applies quantum operations
 
     Returns:
-        (.Operator, pennylane.QNode, .QuantumTape, or Callable): Simplified input.
+        (Operator or MeasurementProcess or qnode (QNode) or quantum function (Callable)
+        or tuple[List[QuantumTape], function]): Simplified input. If an operator or measurement
+        process is provided as input, the simplified input is returned directly. Otherwise, the
+        transformed circuit is returned as described in :func:`qml.transform <pennylane.transform>`.
 
     **Example**
 
@@ -70,11 +72,11 @@ def simplify(input: Union[Operator, MeasurementProcess, QuantumTape, QNode, Call
     Moreover, ``qml.simplify`` can be used to simplify QNodes or quantum functions:
 
     >>> dev = qml.device("default.qubit", wires=2)
-    >>> @qml.simplify
-        @qml.qnode(dev)
-        def circuit():
-            qml.adjoint(qml.prod(qml.RX(1, 0) ** 1, qml.RY(1, 0), qml.RZ(1, 0)))
-            return qml.probs(wires=0)
+    >>> @qml.qnode(dev)
+    ... @qml.simplify
+    ... def circuit():
+    ...     qml.adjoint(qml.prod(qml.RX(1, 0) ** 1, qml.RY(1, 0), qml.RZ(1, 0)))
+    ...     return qml.probs(wires=0)
     >>> circuit()
     tensor([0.64596329, 0.35403671], requires_grad=True)
     >>> list(circuit.tape)
@@ -89,33 +91,21 @@ def simplify(input: Union[Operator, MeasurementProcess, QuantumTape, QNode, Call
             return qml.apply(new_op)
         return input.simplify()
 
-    if isinstance(input, QuantumScript):
-        return input.__class__(
-            [op.simplify() for op in input.operations],
-            [m.simplify() for m in input.measurements],
-            shots=input.shots,
-        )
-
-    if callable(input):
-        old_qfunc = input.func if isinstance(input, QNode) else input
-
-        @wraps(old_qfunc)
-        def qfunc(*args, **kwargs):
-            qs = make_qscript(old_qfunc)(*args, **kwargs)
-            _ = [qml.simplify(op) for op in qs.operations]
-            m = tuple(qml.simplify(m) for m in qs.measurements)
-            return m[0] if len(m) == 1 else m
-
-        if isinstance(input, QNode):
-            return QNode(
-                func=qfunc,
-                device=input.device,
-                interface=input.interface,
-                diff_method=input.diff_method,
-                expansion_strategy=input.expansion_strategy,
-                **input.execute_kwargs,
-                **input.gradient_kwargs,
-            )
-        return qfunc
+    if isinstance(input, QuantumScript) or callable(input):
+        return _simplify_transform(input)
 
     raise ValueError(f"Cannot simplify the object {input} of type {type(input)}.")
+
+
+@qml.transform
+def _simplify_transform(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):
+    with qml.QueuingManager.stop_recording():
+        new_operations = [op.simplify() for op in tape.operations]
+        new_measurements = [m.simplify() for m in tape.measurements]
+
+    new_tape = type(tape)(new_operations, new_measurements, shots=tape.shots)
+
+    def null_processing_fn(res):
+        return res[0]
+
+    return [new_tape], null_processing_fn


Test patch:
diff --git a/tests/ops/functions/test_simplify.py b/tests/ops/functions/test_simplify.py
index bdebb8721..d1f19cabd 100644
--- a/tests/ops/functions/test_simplify.py
+++ b/tests/ops/functions/test_simplify.py
@@ -93,7 +93,7 @@ class TestSimplifyTapes:
             build_op()
 
         tape = QuantumScript.from_queue(q_tape, shots=shots)
-        s_tape = qml.simplify(tape)
+        [s_tape], _ = qml.simplify(tape)
         assert len(s_tape) == 1
         s_op = s_tape[0]
         assert isinstance(s_op, qml.ops.Prod)  # pylint: disable=no-member
@@ -111,7 +111,7 @@ class TestSimplifyTapes:
 
         tape = QuantumScript.from_queue(q_tape)
         simplified_tape_op = qml.PauliZ(1)
-        s_tape = qml.simplify(tape)
+        [s_tape], _ = qml.simplify(tape)
         s_op = s_tape.operations[0]
         assert isinstance(s_op, qml.PauliZ)
         assert s_op.data == simplified_tape_op.data
@@ -138,9 +138,12 @@ class TestSimplifyQNodes:
 
         s_qnode = qml.simplify(qnode)
         assert s_qnode() == qnode()
-        assert len(s_qnode.tape) == 2
-        s_op = s_qnode.tape.operations[0]
-        s_obs = s_qnode.tape.observables[0]
+
+        [s_tape], _ = s_qnode.transform_program([s_qnode.tape])
+        assert len(s_tape) == 2
+
+        s_op = s_tape.operations[0]
+        s_obs = s_tape.observables[0]
         assert isinstance(s_op, qml.PauliZ)
         assert s_op.data == simplified_tape_op.data
         assert s_op.wires == simplified_tape_op.wires