[Capture] unitary_to_rot is plxpr compatible

doc/releases/changelog-dev.md

@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* Added class `qml.capture.transforms.UnitaryToRotInterpreter` that decomposes pennylane operators 
+  following the same API as `qml.transforms.unitary_to_rot` when experimental program capture is enabled.
+  [(#6916)](https://github.com/PennyLaneAI/pennylane/pull/6916)
+
 <h3>Improvements 🛠</h3>
 
 * Add a `qml.capture.pause()` context manager for pausing program capture in an error-safe way.

pennylane/transforms/unitary_to_rot.py

@@ -14,6 +14,7 @@
 """
 A transform for decomposing arbitrary single-qubit QubitUnitary gates into elementary gates.
 """
+from functools import lru_cache, partial
 
 import pennylane as qml
 from pennylane.ops.op_math.decompositions import one_qubit_decomposition, two_qubit_decomposition
@@ -23,7 +24,69 @@
 from pennylane.typing import PostprocessingFn
 
 
-@transform
+@lru_cache
+def _get_plxpr_unitary_to_rot():
+    try:
+        # pylint: disable=import-outside-toplevel
+        from jax import make_jaxpr
+
+        from pennylane.capture import PlxprInterpreter
+        from pennylane.operation import Operator
+    except ImportError:  # pragma: no cover
+        return None, None
+
+    # pylint: disable=redefined-outer-name, too-few-public-methods
+    class UnitaryToRotInterpreter(PlxprInterpreter):
+        """Plxpr Interpreter for applying the ``unitary_to_rot``
+        transform when program capture is enabled."""
+
+        def interpret_operation(self, op: Operator):
+            """Decompose a PennyLane operation instance if it is a QubitUnitary.
+
+            Args:
+                op (Operator): a pennylane operator instance
+
+            Returns:
+                list: The decomposed operations.
+
+            This method is only called when the operator's output is a dropped variable,
+            so the output will not affect later equations in the circuit.
+
+            See also: :meth:`~.interpret_operation_eqn`, :meth:`~.interpret_operation`.
+            """
+            if isinstance(op, qml.QubitUnitary):
+                ops = []
+                with qml.capture.pause():
+                    matrix_shape = qml.math.shape(op.parameters[0])
+                    if matrix_shape == (2, 2):
+                        ops = one_qubit_decomposition(op.parameters[0], op.wires[0])
+                    elif matrix_shape == (4, 4):
+                        ops = two_qubit_decomposition(op.parameters[0], op.wires)
+                # List comprehensions are run in a separate scope.
+                # The automatic insertion of __class__ and self for zero-argument super does not work in such a nested scope.
+                # pylint: disable=super-with-arguments
+                return [super(UnitaryToRotInterpreter, self).interpret_operation(o) for o in ops]
+
+            return super().interpret_operation(op)
+
+    # pylint: disable=redefined-outer-name
+    def unitary_to_rot_plxpr_to_plxpr(
+        jaxpr, consts, _, __, *args
+    ):  # pylint: disable=unused-argument
+        interpreter = UnitaryToRotInterpreter()
+
+        def wrapper(*inner_args):
+            return interpreter.eval(jaxpr, consts, *inner_args)
+
+        return make_jaxpr(wrapper)(*args)
+
+    return UnitaryToRotInterpreter, unitary_to_rot_plxpr_to_plxpr
+
+
+UnitaryToRotInterpreter, unitary_to_rot_plxpr_to_plxpr = _get_plxpr_unitary_to_rot()
+
+
+@partial(transform, plxpr_transform=unitary_to_rot_plxpr_to_plxpr)
 def unitary_to_rot(tape: QuantumScript) -> tuple[QuantumScriptBatch, PostprocessingFn]:
     r"""Quantum function transform to decomposes all instances of single-qubit and
     select instances of two-qubit :class:`~.QubitUnitary` operations to