Multi-level dynamic decompositions

doc/releases/changelog-dev.md

@@ -201,6 +201,7 @@
 * Implemented a `compute_plxpr_decomposition` method in the `qml.operation.Operator` class to apply dynamic decompositions
   with program capture enabled.
   [(#6859)](https://github.com/PennyLaneAI/pennylane/pull/6859)
+  [(#6881)](https://github.com/PennyLaneAI/pennylane/pull/6881)
 
   * Autograph can now be used with custom operations defined outside of the pennylane namespace.
   [(#6931)](https://github.com/PennyLaneAI/pennylane/pull/6931)

pennylane/operation.py

@@ -1361,8 +1361,12 @@ def compute_plxpr_decomposition(*args, **hyperparameters) -> None:
         When the program capture feature is enabled with ``qml.capture.enable()``, the decomposition of the operator
         is computed with this method if it is defined. Otherwise, the :meth:`~.Operator.compute_decomposition` method is used.
 
-        If this method is defined, the control flow operations within the method are recorded in the JAX representation
-        of the operator's decomposition.
+        The exception to this rule is when the operator is returned from the :meth:`~.Operator.compute_decomposition` method
+        of another operator, in which case the decomposition is performed with :meth:`~.Operator.compute_decomposition`
+        (even if this method is defined), and not with this method.
+
+        When ``compute_plxpr_decomposition`` is defined for an operator, the control flow operations within the method
+        (specifying the decomposition of the operator) are recorded in the JAX representation.
 
         This method is experimental and subject to change.
 

pennylane/transforms/decompose.py

@@ -18,9 +18,10 @@
 # pylint: disable=unnecessary-lambda-assignment
 
 import warnings
-from collections.abc import Callable, Generator, Iterable
+from collections.abc import Generator, Iterable
+from copy import copy
 from functools import lru_cache, partial
-from typing import Optional
+from typing import Callable, Optional, Sequence
 
 import pennylane as qml
 from pennylane.transforms.core import transform
@@ -62,12 +63,18 @@
 
 
 @lru_cache
-def _get_plxpr_decompose():  # pylint: disable=missing-docstring
+def _get_plxpr_decompose():  # pylint: disable=missing-docstring, too-many-statements
     try:
         # pylint: disable=import-outside-toplevel
-        from jax import make_jaxpr
+        import jax
+
+        from pennylane.capture.primitives import (
+            cond_prim,
+            ctrl_transform_prim,
+            for_loop_prim,
+            while_loop_prim,
+        )
 
-        from pennylane.capture.primitives import ctrl_transform_prim
     except ImportError:  # pragma: no cover
         return None, None
 
@@ -106,6 +113,7 @@
                 bool: Whether or not ``op`` is valid or needs to be decomposed. ``True`` means
                 that the operator does not need to be decomposed.
             """
+
             if not op.has_decomposition:
                 if not self.gate_set(op):
                     warnings.warn(
@@ -117,15 +125,13 @@
                 return True
             return self.gate_set(op)
 
-        def decompose_operation(self, op: qml.operation.Operator):
+        def decompose_operation(self, op: qml.operation.Operator, current_depth: int = 0):
             """Decompose a PennyLane operation instance if it does not satisfy the
             provided gate set.
 
             Args:
                 op (Operator): a pennylane operator instance
-
-            Returns:
-                Any
+                current_depth (int): the current depth of the decomposition
 
             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.
@@ -135,47 +141,276 @@
             if self.gate_set(op):
                 return self.interpret_operation(op)
 
+            max_expansion = (
+                self.max_expansion - current_depth if self.max_expansion is not None else None
+            )
+
             with qml.capture.pause():
                 decomposition = list(
                     _operator_decomposition_gen(
-                        op, self.stopping_condition, max_expansion=self.max_expansion
+                        op,
+                        self.stopping_condition,
+                        max_expansion=max_expansion,
                     )
                 )
 
             return [self.interpret_operation(decomp_op) for decomp_op in decomposition]
 
-        def interpret_operation_eqn(self, eqn):
+        def _evaluate_jaxpr_decomposition(self, op: qml.operation.Operator, current_depth: int = 0):
+            """Creates and evaluates a Jaxpr of the plxpr decomposition of an operator."""
+
+            if self.gate_set(op):
+                return self.interpret_operation(op)
+
+            if self.max_expansion is not None and current_depth >= self.max_expansion:
+                return self.interpret_operation(op)
+
+            args = (*op.parameters, *op.wires)
+            jaxpr_decomp = qml.capture.make_plxpr(
+                partial(op.compute_plxpr_decomposition, **op.hyperparameters)
+            )(*args)
+            current_depth += 1
+
+            return self.eval(
+                jaxpr_decomp.jaxpr, jaxpr_decomp.consts, *args, current_depth=current_depth
+            )
+
+        def eval(
+            self, jaxpr: "jax.core.Jaxpr", consts: Sequence, *args, current_depth: int = 0
+        ) -> list:
+            """
+            Evaluates a jaxpr, which can also be generated by a dynamic decomposition.
+
+            Args:
+                jaxpr_decomp (jax.core.Jaxpr): the Jaxpr to evaluate
+                consts (list[TensorLike]): the constant variables for the jaxpr
+                *args: the arguments to use in the evaluation
+                current_depth (int): the current depth of the decomposition
+            """
+
+            # We update the 'self._env' environment directly because the jaxpr of the decomposition
+            # can be called while evaluating another jaxpr of the previous decomposition.
+
+            for arg, invar in zip(args, jaxpr.invars, strict=True):
+                self._env[invar] = arg
+            for const, constvar in zip(consts, jaxpr.constvars, strict=True):
+                self._env[constvar] = const
+
+            for eq in jaxpr.eqns:
+
+                prim_type = getattr(eq.primitive, "prim_type", "")
+                custom_handler = self._primitive_registrations.get(eq.primitive, None)
+
+                if custom_handler:
+
+                    invals = [self.read(invar) for invar in eq.invars]
+                    control_flow_handlers = {"handle_for_loop", "handle_while_loop", "handle_cond"}
+                    extra_kwargs = (
+                        {"current_depth": current_depth}
+                        if custom_handler.__name__ in control_flow_handlers
+                        else {}
+                    )
+                    outvals = custom_handler(self, *invals, **eq.params, **extra_kwargs)
+
+                elif prim_type == "operator":
+                    outvals = self.interpret_operation_eqn(eq, current_depth)
+                elif prim_type == "measurement":
+                    outvals = self.interpret_measurement_eqn(eq)
+                else:
+                    invals = [self.read(invar) for invar in eq.invars]
+                    subfuns, params = eq.primitive.get_bind_params(eq.params)
+                    outvals = eq.primitive.bind(*subfuns, *invals, **params)
+
+                if not eq.primitive.multiple_results:
+                    outvals = [outvals]
+
+                for outvar, outval in zip(eq.outvars, outvals, strict=True):
+                    self._env[outvar] = outval
+
+            outvals = []
+            for var in jaxpr.outvars:
+                outval = self.read(var)
+                if isinstance(outval, qml.operation.Operator):
+                    outvals.append(self.interpret_operation(outval))
+                else:
+                    outvals.append(outval)
+
+            return outvals
+
+        def interpret_operation_eqn(self, eqn, current_depth: int = 0):
             """Interpret an equation corresponding to an operator.
 
             Args:
                 eqn (jax.core.JaxprEqn): a jax equation for an operator.
+                current_depth (int): the current depth of the decomposition.
 
             See also: :meth:`~.interpret_operation`.
 
             """
             invals = (self.read(invar) for invar in eqn.invars)
+
             with qml.QueuingManager.stop_recording():
                 op = eqn.primitive.impl(*invals, **eqn.params)
-            if eqn.outvars[0].__class__.__name__ == "DropVar":
 
-                if op.has_plxpr_decomposition:
+            if not eqn.outvars[0].__class__.__name__ == "DropVar":
+                return op
 
-                    args = (*op.parameters, *op.wires)
-                    qml.capture.run_autograph(op.compute_plxpr_decomposition)(
-                        *args, **op.hyperparameters
-                    )
+            if not op.has_plxpr_decomposition:
+                return self.decompose_operation(op, current_depth)
 
-                else:
+            return self._evaluate_jaxpr_decomposition(op, current_depth)
 
-                    return self.decompose_operation(op)
+    def jaxpr_to_jaxpr_decomp(
+        interpreter: DecomposeInterpreter, jaxpr: "jax.core.Jaxpr", consts, *args, current_depth
+    ) -> "jax.core.Jaxpr":
 
-            return op
+        f = partial(interpreter.eval, jaxpr, consts, current_depth=current_depth)
+
+        return jax.make_jaxpr(f)(*args)
 
     # pylint: disable=unused-variable,missing-function-docstring
     @DecomposeInterpreter.register_primitive(ctrl_transform_prim)
     def handle_ctrl_transform(*_, **__):
         raise NotImplementedError
 
+    # We register the primitives to propagate the current depth
+    # in the dynamic decomposition evaluation with program capture enabled.
+
+    @DecomposeInterpreter.register_primitive(cond_prim)
+    def handle_cond(self, *invals, jaxpr_branches, consts_slices, args_slice, current_depth=0):
+        """Handle a cond primitive."""
+
+        args = invals[args_slice]
+
+        new_jaxprs = []
+        new_consts = []
+        new_consts_slices = []
+        end_const_ind = len(jaxpr_branches)
+
+        for const_slice, jaxpr in zip(consts_slices, jaxpr_branches):
+            consts = invals[const_slice]
+            if jaxpr is None:
+                new_jaxprs.append(None)
+                new_consts_slices.append(slice(0, 0))
+            else:
+                new_jaxpr = jaxpr_to_jaxpr_decomp(
+                    copy(self), jaxpr, consts, *args, current_depth=current_depth
+                )
+                new_jaxprs.append(new_jaxpr.jaxpr)
+                new_consts.extend(new_jaxpr.consts)
+                new_consts_slices.append(
+                    slice(end_const_ind, end_const_ind + len(new_jaxpr.consts))
+                )
+                end_const_ind += len(new_jaxpr.consts)
+
+        new_args_slice = slice(end_const_ind, None)
+        return cond_prim.bind(
+            *invals[: len(jaxpr_branches)],
+            *new_consts,
+            *args,
+            jaxpr_branches=new_jaxprs,
+            consts_slices=new_consts_slices,
+            args_slice=new_args_slice,
+        )
+
+    @DecomposeInterpreter.register_primitive(for_loop_prim)
+    def handle_for_loop(
+        self,
+        start,
+        stop,
+        step,
+        *args,
+        jaxpr_body_fn,
+        consts_slice,
+        args_slice,
+        abstract_shapes_slice,
+        current_depth=0,
+    ):
+        """Handle a for loop primitive."""
+
+        consts = args[consts_slice]
+        init_state = args[args_slice]
+        abstract_shapes = args[abstract_shapes_slice]
+
+        new_jaxpr_body_fn = jaxpr_to_jaxpr_decomp(
+            copy(self),
+            jaxpr_body_fn,
+            consts,
+            *abstract_shapes,
+            start,
+            *init_state,
+            current_depth=current_depth,
+        )
+
+        consts_slice = slice(0, len(new_jaxpr_body_fn.consts))
+        abstract_shapes_slice = slice(consts_slice.stop, consts_slice.stop + len(abstract_shapes))
+        args_slice = slice(abstract_shapes_slice.stop, None)
+        return for_loop_prim.bind(
+            start,
+            stop,
+            step,
+            *new_jaxpr_body_fn.consts,
+            *abstract_shapes,
+            *init_state,
+            jaxpr_body_fn=new_jaxpr_body_fn.jaxpr,
+            consts_slice=consts_slice,
+            args_slice=args_slice,
+            abstract_shapes_slice=abstract_shapes_slice,
+        )
+
+    @DecomposeInterpreter.register_primitive(while_loop_prim)
+    def handle_while_loop(
+        self,
+        *invals,
+        jaxpr_body_fn,
+        jaxpr_cond_fn,
+        body_slice,
+        cond_slice,
+        args_slice,
+        abstract_shapes_slice,
+        current_depth=0,
+    ):
+        """Handle a while loop primitive."""
+        consts_body = invals[body_slice]
+        consts_cond = invals[cond_slice]
+        init_state = invals[args_slice]
+        abstract_shapes = invals[abstract_shapes_slice]
+
+        new_jaxpr_body_fn = jaxpr_to_jaxpr_decomp(
+            copy(self),
+            jaxpr_body_fn,
+            consts_body,
+            *abstract_shapes,
+            *init_state,
+            current_depth=current_depth,
+        )
+        new_jaxpr_cond_fn = jaxpr_to_jaxpr_decomp(
+            copy(self),
+            jaxpr_cond_fn,
+            consts_cond,
+            *abstract_shapes,
+            *init_state,
+            current_depth=current_depth,
+        )
+
+        body_consts = slice(0, len(new_jaxpr_body_fn.consts))
+        cond_consts = slice(body_consts.stop, body_consts.stop + len(new_jaxpr_cond_fn.consts))
+        abstract_shapes_slice = slice(cond_consts.stop, cond_consts.stop + len(abstract_shapes))
+        args_slice = slice(abstract_shapes_slice.stop, None)
+
+        return while_loop_prim.bind(
+            *new_jaxpr_body_fn.consts,
+            *new_jaxpr_cond_fn.consts,
+            *abstract_shapes,
+            *init_state,
+            jaxpr_body_fn=new_jaxpr_body_fn.jaxpr,
+            jaxpr_cond_fn=new_jaxpr_cond_fn.jaxpr,
+            body_slice=body_consts,
+            cond_slice=cond_consts,
+            args_slice=args_slice,
+            abstract_shapes_slice=abstract_shapes_slice,
+        )
+
     def decompose_plxpr_to_plxpr(
         jaxpr, consts, targs, tkwargs, *args
     ):  # pylint: disable=unused-argument
@@ -187,7 +422,7 @@
         def wrapper(*inner_args):
             return decomposer.eval(jaxpr, consts, *inner_args)
 
-        return make_jaxpr(wrapper)(*args)
+        return jax.make_jaxpr(wrapper)(*args)
 
     return DecomposeInterpreter, decompose_plxpr_to_plxpr