Exploring how to handle dynamic decompositions with PLxPR enabled

pennylane/transforms/decompose.py

@@ -20,7 +20,7 @@
 import warnings
 from collections.abc import Callable, Generator, Iterable
 from functools import lru_cache, partial
-from typing import Optional
+from typing import Optional, Sequence
 
 import pennylane as qml
 from pennylane.transforms.core import transform
@@ -187,6 +187,92 @@
 DecomposeInterpreter, decompose_plxpr_to_plxpr = _get_plxpr_decompose()
 
 
+@lru_cache
+def _get_plxpr_dynamic_decompose():  # pylint: disable=missing-docstring
+    try:
+        # pylint: disable=import-outside-toplevel
+        # pylint: disable=unused-import
+        import jax
+        from pennylane.capture.primitives import AbstractMeasurement, AbstractOperator
+    except ImportError:  # pragma: no cover
+        return None, None
+
+    # pylint: disable=redefined-outer-name
+
+    class DynamicDecomposeInterpreter(qml.capture.PlxprInterpreter):
+        """
+        Experimental Plxpr Interpreter for applying a dynamic decomposition to operations program capture is enabled.
+
+        """
+
+        def eval_dynamic_decomposition(
+            self, jaxpr_decomp: "jax.core.Jaxpr", consts: Sequence, *args
+        ):
+            """
+            Evaluate a dynamic decomposition of a Jaxpr.
+
+            Args:
+                jaxpr_decomp (jax.core.Jaxpr): the Jaxpr to evaluate
+                consts (Sequence): the constants to use in the evaluation
+                *args: the arguments to use in the evaluation
+
+            """
+
+            for arg, invar in zip(args, jaxpr_decomp.invars, strict=True):
+                self._env[invar] = arg
+            for const, constvar in zip(consts, jaxpr_decomp.constvars, strict=True):
+                self._env[constvar] = const
+
+            for inner_eqn in jaxpr_decomp.eqns:
+
+                custom_handler = self._primitive_registrations.get(inner_eqn.primitive, None)
+
+                if custom_handler:
+                    invals = [self.read(invar) for invar in inner_eqn.invars]
+                    outvals = custom_handler(self, *invals, **inner_eqn.params)
+
+                elif isinstance(inner_eqn.outvars[0].aval, AbstractOperator):
+                    outvals = super().interpret_operation_eqn(inner_eqn)
+                elif isinstance(inner_eqn.outvars[0].aval, AbstractMeasurement):
+                    outvals = super().interpret_measurement_eqn(inner_eqn)
+                else:
+                    invals = [self.read(invar) for invar in inner_eqn.invars]
+                    outvals = inner_eqn.primitive.bind(*invals, **inner_eqn.params)
+
+                if not inner_eqn.primitive.multiple_results:
+                    outvals = [outvals]
+
+                for inner_outvar, inner_outval in zip(inner_eqn.outvars, outvals, strict=True):
+                    self._env[inner_outvar] = inner_outval
+
+        def interpret_operation_eqn(self, eqn: "jax.core.JaxprEqn"):
+            """
+            Interpret an equation corresponding to an operator.
+
+            Args:
+                eqn (jax.core.JaxprEqn): a jax equation for an operator.
+            """
+
+            invals = (self.read(invar) for invar in eqn.invars)
+            with qml.QueuingManager.stop_recording():
+                op = eqn.primitive.impl(*invals, **eqn.params)
+
+            if hasattr(op, "_compute_plxpr_decomposition"):
+
+                jaxpr_decomp = op._plxpr_decomposition()
+                args = (*op.parameters, *op.wires, *op.hyperparameters)
+                return self.eval_dynamic_decomposition(
+                    jaxpr_decomp.jaxpr, jaxpr_decomp.consts, *args
+                )
+
+            return super().interpret_operation_eqn(eqn)
+
+    return DynamicDecomposeInterpreter
+
+
+DynamicDecomposeInterpreter = _get_plxpr_dynamic_decompose()
+
+
 @partial(transform, plxpr_transform=decompose_plxpr_to_plxpr)
 def decompose(tape, gate_set=None, max_expansion=None):
     """Decomposes a quantum circuit into a user-specified gate set.

tests/capture/transforms/test_capture_decompose.py

@@ -30,6 +30,7 @@
 )
 from pennylane.transforms.decompose import DecomposeInterpreter, decompose_plxpr_to_plxpr
 
+
 pytestmark = [pytest.mark.jax, pytest.mark.usefixtures("enable_disable_plxpr")]