Defer calculating classical cotransform and argnums till needed

doc/releases/changelog-dev.md

@@ -395,6 +395,11 @@ such as `shots`, `rng` and `prng_key`.
 * Moved all interface handling logic to `interface_utils.py` in the `qml.math` module.
   [(#6649)](https://github.com/PennyLaneAI/pennylane/pull/6649)
 
+* `qml.execute` can now be used with `diff_method="best"`.
+  Classical cotransform information is now handled lazily by the workflow. Gradient method
+  validation and program setup is now handled inside of `qml.execute`, instead of in `QNode`.
+  [(#6716)](https://github.com/PennyLaneAI/pennylane/pull/6716)
+
 * Added PyTree support for measurements in a circuit. 
   [(#6378)](https://github.com/PennyLaneAI/pennylane/pull/6378)
 

pennylane/math/interface_utils.py

@@ -209,7 +209,7 @@ def get_deep_interface(value):
     return _get_interface_of_single_tensor(itr)
 
 
-def get_canonical_interface_name(user_input: Union[str, Interface]) -> Interface:
+def get_canonical_interface_name(user_input: Union[str, Interface, None]) -> Interface:
     """Helper function to get the canonical interface name.
 
     Args:

pennylane/transforms/core/transform_program.py

@@ -14,6 +14,7 @@
 """
 This module contains the ``TransformProgram`` class.
 """
+from collections import namedtuple
 from collections.abc import Sequence
 from functools import partial
 from typing import Optional, overload
@@ -24,6 +25,20 @@
 
 from .transform_dispatcher import TransformContainer, TransformDispatcher, TransformError
 
+CotransfromCache = namedtuple("CotransformCache", ("qnode", "args", "kwargs"))
+
+
+def _get_interface(qnode, args, kwargs):
+    if qnode.interface == "auto":
+        interface = qml.math.get_interface(*args, *list(kwargs.values()))
+        try:
+            interface = qml.math.get_canonical_interface_name(interface).value
+        except ValueError:
+            interface = "numpy"
+    else:
+        interface = qnode.interface
+    return interface
+
 
 def _numpy_jac(*_, **__) -> qml.typing.TensorLike:
     raise qml.QuantumFunctionError("No trainable parameters.")
@@ -132,7 +147,7 @@ def _batch_postprocessing(
 
     Keyword Args:
         individual_fns (List[Callable]): postprocessing functions converting a batch of results into a single result
-           corresponding to only a single :class:`~.QuantumTape`.
+            corresponding to only a single :class:`~.QuantumTape`.
         slices (List[slice]): the indices for the results that correspond to each individual post processing function.
 
     >>> results = (1.0, 2.0, 3.0, 4.0)
@@ -200,6 +215,12 @@ class TransformProgram:
 
     The order of execution is the order in the list containing the containers.
 
+    Args:
+        initial_program (Optional[Sequence[TransformContainer]]): A sequence of transforms to
+            initialize the program with
+        cotransform_cache (Optional[CotransformCache]): A named tuple containing the ``qnode``,
+            ``args``, and ``kwargs`` required to compute classical cotransforms.
+
     The main case where one would have to interact directly with a transform program is when developing a
     :class:`Device <pennylane.devices.Device>`. In this case, the pre-processing method of a device
     returns a transform program. You should directly refer to the device API documentation for more details.
@@ -243,10 +264,13 @@ class TransformProgram:
 
     """
 
-    def __init__(self, initial_program: Optional[Sequence] = None):
+    def __init__(
+        self,
+        initial_program: Optional[Sequence[TransformContainer]] = None,
+        cotransform_cache: Optional[CotransfromCache] = None,
+    ):
         self._transform_program = list(initial_program) if initial_program else []
-        self._classical_jacobians = None
-        self._argnums = None
+        self.cotransform_cache = cotransform_cache
 
     def __iter__(self):
         """list[TransformContainer]: Return an iterator to the underlying transform program."""
@@ -270,15 +294,22 @@ def __getitem__(self, idx):
     def __bool__(self) -> bool:
         return bool(self._transform_program)
 
-    def __add__(self, other):
+    def __add__(self, other: "TransformProgram") -> "TransformProgram":
         if self.has_final_transform and other.has_final_transform:
             raise TransformError("The transform program already has a terminal transform.")
 
         transforms = self._transform_program + other._transform_program
         if self.has_final_transform:
             transforms.append(transforms.pop(len(self) - 1))
 
-        return TransformProgram(transforms)
+        cotransform_cache = None
+        if self.cotransform_cache:
+            if other.cotransform_cache:
+                raise ValueError("Cannot add two transform programs with cotransform caches.")
+            cotransform_cache = self.cotransform_cache
+        elif other.cotransform_cache:
+            cotransform_cache = other.cotransform_cache
+        return TransformProgram(transforms, cotransform_cache=cotransform_cache)
 
     def __repr__(self):
         """The string representation of the transform program class."""
@@ -444,17 +475,9 @@ def has_classical_cotransform(self) -> bool:
 
     def set_classical_component(self, qnode, args, kwargs):
         """Set the classical jacobians and argnums if the transform is hybrid with a classical cotransform."""
-        if not self.has_classical_cotransform():
-            return
-        hybrid = self[-1].kwargs.pop("hybrid", True)  # pylint: disable=no-member
-
-        if hybrid:
-            argnums = self[-1].kwargs.pop("argnums", None)  # pylint: disable=no-member
-            self._classical_jacobians = [
-                self._get_classical_jacobian(index, qnode, args, kwargs, argnums)
-                for index, _ in enumerate(self)
-            ]
-            self._set_all_argnums(qnode, args, kwargs, argnums)
+        # pylint: disable=no-member
+        if self.has_classical_cotransform() and self[-1].kwargs.get("hybrid", True):
+            self.cotransform_cache = CotransfromCache(qnode, args, kwargs)
 
     def prune_dynamic_transform(self, type_to_keep=1):
         """Ensures that only one or none ``dynamic_one_shot`` is applied.
@@ -484,16 +507,20 @@ def prune_dynamic_transform(self, type_to_keep=1):
         return found
 
     # pylint: disable=too-many-arguments, too-many-positional-arguments
-    def _get_classical_jacobian(self, index: int, qnode, args, kwargs, argnums):
-        if not self[index].classical_cotransform:
+    def _get_classical_jacobian(self, index: int):
+        if self.cotransform_cache is None or not self[index].classical_cotransform:
             return None
-        if qnode.interface == "jax" and "argnum" in self[index].kwargs:
+        argnums = self[-1].kwargs.get("argnums", None)  # pylint: disable=no-member
+        qnode, args, kwargs = self.cotransform_cache
+
+        interface = _get_interface(qnode, args, kwargs)
+        if interface == "jax" and "argnum" in self[index].kwargs:
             raise qml.QuantumFunctionError(
                 "argnum does not work with the Jax interface. You should use argnums instead."
             )
 
         f = partial(_classical_preprocessing, qnode, self[:index])
-        classical_jacobian = _jac_map[qnode.interface](f, argnums, *args, **kwargs)
+        classical_jacobian = _jac_map[interface](f, argnums, *args, **kwargs)
 
         # autograd and tf cant handle pytrees, so need to unsqueeze the squeezing
         # done in _classical_preprocessing
@@ -504,22 +531,22 @@ def _get_classical_jacobian(self, index: int, qnode, args, kwargs, argnums):
             classical_jacobian = [classical_jacobian]
         return classical_jacobian
 
-    def _set_all_argnums(self, qnode, args, kwargs, argnums):
+    def _get_argnums(self, index):
         """It can be used inside the QNode to set all argnums (tape level) using argnums from the argnums at the QNode
         level.
         """
-
-        argnums_list = []
-        for index, transform in enumerate(self):
-            argnums = [0] if qnode.interface in ["jax", "jax-jit"] and argnums is None else argnums
-            # pylint: disable=protected-access
-            if (transform._use_argnum or transform.classical_cotransform) and argnums:
-                params = _jax_argnums_to_tape_trainable(qnode, argnums, self[:index], args, kwargs)
-                argnums_list.append([qml.math.get_trainable_indices(param) for param in params])
-            else:
-                argnums_list.append(None)
-
-        self._argnums = argnums_list
+        if self.cotransform_cache is None:
+            return None
+        qnode, args, kwargs = self.cotransform_cache
+        interface = _get_interface(qnode, args, kwargs)
+        transform = self[index]
+        argnums = self[-1].kwargs.get("argnums", None)  # pylint: disable=no-member
+        argnums = [0] if interface in ["jax", "jax-jit"] and argnums is None else argnums
+        # pylint: disable=protected-access
+        if (transform._use_argnum or transform.classical_cotransform) and argnums:
+            params = _jax_argnums_to_tape_trainable(qnode, argnums, self[:index], args, kwargs)
+            return [qml.math.get_trainable_indices(param) for param in params]
+        return None
 
     def __call__(
         self, tapes: QuantumScriptBatch
@@ -531,7 +558,9 @@ def __call__(
 
         for i, transform_container in enumerate(self):
             transform, targs, tkwargs, cotransform, _, _, _ = transform_container
-
+            tkwargs = {
+                key: value for key, value in tkwargs.items() if key not in {"argnums", "hybrid"}
+            }
             execution_tapes = []
             fns = []
             slices = []
@@ -541,9 +570,11 @@ def __call__(
 
             start = 0
             start_classical = 0
+            classical_jacobians = self._get_classical_jacobian(i)
+            argnums = self._get_argnums(i)
             for j, tape in enumerate(tapes):
-                if self._argnums is not None and self._argnums[i] is not None:
-                    tape.trainable_params = self._argnums[i][j]
+                if argnums is not None:
+                    tape.trainable_params = argnums[j]
                 new_tapes, fn = transform(tape, *targs, **tkwargs)
                 execution_tapes.extend(new_tapes)
 
@@ -552,14 +583,14 @@ def __call__(
                 slices.append(slice(start, end))
                 start = end
 
-                if cotransform and self._classical_jacobians:
+                if cotransform and classical_jacobians:
                     classical_fns.append(
-                        partial(cotransform, cjac=self._classical_jacobians[i][j], tape=tape)
+                        partial(cotransform, cjac=classical_jacobians[j], tape=tape)
                     )
                     slices_classical.append(slice(start_classical, start_classical + 1))
                     start_classical += 1
 
-            if cotransform and self._classical_jacobians:
+            if cotransform and classical_jacobians:
                 batch_postprocessing_classical = partial(
                     _batch_postprocessing, individual_fns=classical_fns, slices=slices_classical
                 )
@@ -581,5 +612,4 @@ def __call__(
         postprocessing_fn.__doc__ = _apply_postprocessing_stack.__doc__
 
         # Reset classical jacobians
-        self._classical_jacobians = []
         return tuple(tapes), postprocessing_fn

pennylane/workflow/_setup_transform_program.py

@@ -87,7 +87,9 @@ def _setup_transform_program(
 
     device_transform_program = device.preprocess_transforms(resolved_execution_config)
 
-    full_transform_program = qml.transforms.core.TransformProgram(user_transform_program)
+    full_transform_program = qml.transforms.core.TransformProgram(
+        user_transform_program, cotransform_cache=user_transform_program.cotransform_cache
+    )
     inner_transform_program = qml.transforms.core.TransformProgram()
 
     # Add the gradient expand to the program if necessary

pennylane/workflow/execution.py

@@ -31,7 +31,7 @@
 from pennylane.typing import ResultBatch
 
 from ._setup_transform_program import _setup_transform_program
-from .resolution import _resolve_interface
+from .resolution import _resolve_execution_config, _resolve_interface
 from .run import run
 
 logger = logging.getLogger(__name__)
@@ -39,7 +39,7 @@
 
 
 # pylint: disable=too-many-arguments
 def execute(
     tapes: QuantumScriptBatch,
     device: Union["qml.devices.LegacyDevice", "qml.devices.Device"],
     diff_method: Optional[Union[Callable, str, qml.transforms.core.TransformDispatcher]] = None,
@@ -187,21 +187,13 @@
             "::L".join(str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]),
         )
 
+    if not tapes:
+        return ()
+
     ### Specifying and preprocessing variables ####
 
     interface = _resolve_interface(interface, tapes)
     # Only need to calculate derivatives with jax when we know it will be executed later.
-    if interface in {Interface.JAX, Interface.JAX_JIT}:
-        grad_on_execution = grad_on_execution if isinstance(diff_method, Callable) else False
-
-    if (
-        device_vjp
-        and isinstance(device, qml.devices.LegacyDeviceFacade)
-        and "lightning" not in getattr(device, "short_name", "").lower()
-    ):
-        raise qml.QuantumFunctionError(
-            "device provided jacobian products are not compatible with the old device interface."
-        )
 
     gradient_kwargs = gradient_kwargs or {}
     mcm_config = mcm_config or {}
@@ -215,7 +207,9 @@
             gradient_keyword_arguments=gradient_kwargs,
             derivative_order=max_diff,
         )
-        config = device.setup_execution_config(config)
+        config = _resolve_execution_config(
+            config, device, tapes, transform_program=transform_program
+        )
 
     config = replace(
         config,
@@ -224,6 +218,7 @@
     )
 
     if transform_program is None or inner_transform is None:
+        transform_program = transform_program or qml.transforms.core.TransformProgram()
         transform_program, inner_transform = _setup_transform_program(
             transform_program, device, config, cache, cachesize
         )