Move device preprocessing to inner execute

doc/releases/changelog-dev.md

@@ -209,6 +209,9 @@
   [(#5758)](https://github.com/PennyLaneAI/pennylane/pull/5758/)
   [(#5638)](https://github.com/PennyLaneAI/pennylane/pull/5638/)
 
+* Device preprocess transforms now happen inside the ml boundary.
+  [(#5791)](https://github.com/PennyLaneAI/pennylane/pull/5791)
+
 * `qml.qchem.molecular_dipole` function is added for calculating the dipole operator using "dhf" and "openfermion" backends.
   [(#5764)](https://github.com/PennyLaneAI/pennylane/pull/5764)
 

pennylane/devices/default_qubit.py

@@ -502,7 +502,10 @@ def preprocess(
 
         transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
         transform_program.add_transform(
-            mid_circuit_measurements, device=self, mcm_config=config.mcm_config
+            mid_circuit_measurements,
+            device=self,
+            mcm_config=config.mcm_config,
+            interface=config.interface,
         )
         transform_program.add_transform(
             decompose,

pennylane/devices/preprocess.py

@@ -147,7 +147,10 @@ def validate_device_wires(
 
 @transform
 def mid_circuit_measurements(
-    tape: qml.tape.QuantumTape, device, mcm_config=MCMConfig()
+    tape: qml.tape.QuantumTape,
+    device,
+    mcm_config=MCMConfig(),
+    interface=None,
 ) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Provide the transform to handle mid-circuit measurements.
 
@@ -162,7 +165,7 @@ def mid_circuit_measurements(
         mcm_method = "one-shot" if tape.shots else "deferred"
 
     if mcm_method == "one-shot":
-        return qml.dynamic_one_shot(tape)
+        return qml.dynamic_one_shot(tape, interface=interface)
     return qml.defer_measurements(tape, device=device)
 
 

pennylane/ops/op_math/condition.py

@@ -28,14 +28,14 @@ class ConditionalTransformError(ValueError):
     """Error for using qml.cond incorrectly"""
 
 
-class Conditional(SymbolicOp):
+class Conditional(SymbolicOp, Operation):
     """A Conditional Operation.
 
     Unless you are a Pennylane plugin developer, **you should NOT directly use this class**,
     instead, use the :func:`qml.cond <.cond>` function.
 
     The ``Conditional`` class is a container class that defines an operation
-    that should by applied relative to a single measurement value.
+    that should be applied relative to a single measurement value.
 
     Support for executing ``Conditional`` operations is device-dependent. If a
     device doesn't support mid-circuit measurements natively, then the QNode
@@ -54,13 +54,15 @@ def __init__(self, expr, then_op: Type[Operation], id=None):
         self.hyperparameters["meas_val"] = expr
         self._name = f"Conditional({then_op.name})"
         super().__init__(then_op, id=id)
+        if self.grad_recipe is None:
+            self.grad_recipe = [None] * self.num_params
 
     def label(self, decimals=None, base_label=None, cache=None):
         return self.base.label(decimals=decimals, base_label=base_label, cache=cache)
 
     @property
     def meas_val(self):
-        "the measurement outcome value to consider from `expr` argument"
+        """the measurement outcome value to consider from `expr` argument"""
         return self.hyperparameters["meas_val"]
 
     @property

pennylane/transforms/core/transform_program.py

@@ -17,8 +17,6 @@
 from functools import partial
 from typing import Callable, List, Optional, Sequence, Tuple, Union
 
-import numpy as np
-
 import pennylane as qml
 from pennylane.tape import QuantumTape
 from pennylane.typing import Result, ResultBatch
@@ -354,24 +352,32 @@ def set_classical_component(self, qnode, args, kwargs):
             self._set_all_classical_jacobians(qnode, args, kwargs, argnums)
             self._set_all_argnums(qnode, args, kwargs, argnums)
 
-    def prune_dynamic_transform(self):
-        """Ensure a single ``dynamic_one_shot`` transform is applied."""
-        trans_type = np.zeros(len(self._transform_program), dtype=np.int32)
-        for i, t in enumerate(self._transform_program):
-            if "dynamic_one_shot" in str(t):
-                trans_type[i] = 1
-            if "mid_circuit_measurements" in str(t):
-                trans_type[i] = 2
-        if sum(trans_type) < 2:
-            return
-        keep = 2 if 2 in trans_type else 1
+    def prune_dynamic_transform(self, type_to_keep=1):
+        """Ensures that only one or none ``dynamic_one_shot`` is applied.
+
+        Args:
+            type_to_keep (int): The type of the dynamic transform to keep. 0: keep none,
+                1: dynamic_one_shot or mid_circuit_measurements, 2: only mid_circuit_measurements.
+
+        Returns:
+            bool: ``True`` if a dynamic transform was found, ``False`` otherwise.
+
+        """
+
+        i = len(self._transform_program) - 1
         found = False
-        for i, ttype in enumerate(reversed(trans_type)):
-            if not found and ttype == keep:
+        while i >= 0:
+            t = self._transform_program[i]
+            if "mid_circuit_measurements" in str(t) and type_to_keep > 0:
+                type_to_keep = 0  # keep this and do not keep the rest
+                found = True
+            elif "dynamic_one_shot" in str(t) and type_to_keep == 1:
+                type_to_keep = 0  # keep this and do not keep the rest
                 found = True
-                continue
-            if found and ttype in [1, 2]:
-                self._transform_program.pop(len(self._transform_program) - 1 - i)
+            elif "dynamic_one_shot" in str(t) or "mid_circuit_measurements" in str(t):
+                self._transform_program.pop(i)
+            i -= 1
+        return found
 
     def _set_all_classical_jacobians(
         self, qnode, args, kwargs, argnums

pennylane/transforms/dynamic_one_shot.py

@@ -118,6 +118,8 @@ def func(x, y):
 
     aux_tapes = [init_auxiliary_tape(t) for t in tapes]
 
+    interface = kwargs.get("interface", None)
+
     def processing_fn(results, has_partitioned_shots=None, batched_results=None):
         if batched_results is None and batch_size is not None:
             # If broadcasting, recursively process the results for each batch. For each batch
@@ -141,7 +143,7 @@ def processing_fn(results, has_partitioned_shots=None, batched_results=None):
             return tuple(final_results)
         if not tape.shots.has_partitioned_shots:
             results = results[0]
-        return parse_native_mid_circuit_measurements(tape, aux_tapes, results)
+        return parse_native_mid_circuit_measurements(tape, aux_tapes, results, interface=interface)
 
     return aux_tapes, processing_fn
 
@@ -208,7 +210,10 @@ def init_auxiliary_tape(circuit: qml.tape.QuantumScript):
 
 
 def parse_native_mid_circuit_measurements(
-    circuit: qml.tape.QuantumScript, aux_tapes: qml.tape.QuantumScript, results: TensorLike
+    circuit: qml.tape.QuantumScript,
+    aux_tapes: qml.tape.QuantumScript,
+    results: TensorLike,
+    interface=None,
 ):
     """Combines, gathers and normalizes the results of native mid-circuit measurement runs.
 
@@ -228,7 +233,7 @@ def measurement_with_no_shots(measurement):
             else np.nan
         )
 
-    interface = qml.math.get_deep_interface(circuit.data)
+    interface = interface or qml.math.get_deep_interface(circuit.data)
     interface = "numpy" if interface == "builtins" else interface
 
     all_mcms = [op for op in aux_tapes[0].operations if is_mcm(op)]

pennylane/workflow/execution.py

@@ -253,13 +253,19 @@ def device_expansion_function(tape):  # pylint: disable=function-redefined
 
 
 def _make_inner_execute(
-    device, override_shots, cache, expand_fn=None, execution_config=None, numpy_only=True
+    device,
+    override_shots,
+    cache,
+    inner_transform,
+    expand_fn=None,
+    execution_config=None,
+    numpy_only=True,
 ) -> Callable:
     """Construct the function that will execute the tapes inside the ml framework registration
     for the 1st order derivatives.
 
     Steps in between the ml framework execution and the device are:
-    - device expansion (old device)
+    - device expansion (old device) or device preprocessing (new device)
     - conversion to numpy
     - caching
 
@@ -287,23 +293,24 @@ def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch:
 
         Closure Variables:
             expand_fn (Callable[[QuantumTape], QuantumTape]): A device preprocessing step
-            numpy_only (bool): whether or not to convert the data to numpy or leave as is
+            numpy_only (bool): whether to convert the data to numpy or leave as is
             device_execution (Callable[[Sequence[QuantumTape]], ResultBatch])
             cache (None | MutableMapping): The cache to use. If ``None``, caching will not occur.
         """
-        transform_program = qml.transforms.core.TransformProgram()
+
+        transform_program = qml.transforms.core.TransformProgram(inner_transform)
 
         if numpy_only:
             transform_program.add_transform(qml.transforms.convert_to_numpy_parameters)
 
         if cache is not None:
             transform_program.add_transform(_cache_transform, cache=cache)
 
+        transformed_tapes, transform_post_processing = transform_program(tapes)
+
         # TODO: Apply expand_fn() as transform.
         if expand_fn:
-            tapes = tuple(expand_fn(t) for t in tapes)
-
-        transformed_tapes, transform_post_processing = transform_program(tapes)
+            transformed_tapes = tuple(expand_fn(t) for t in transformed_tapes)
 
         if transformed_tapes:
             results = device_execution(transformed_tapes)
@@ -407,6 +414,7 @@ def execute(
     gradient_fn: Optional[Union[Callable, str]] = None,
     interface="auto",
     transform_program=None,
+    inner_transform=None,
     config=None,
     grad_on_execution="best",
     gradient_kwargs=None,
@@ -435,6 +443,7 @@ def execute(
             This affects the types of parameters that can exist on the input tapes.
             Available options include ``autograd``, ``torch``, ``tf``, ``jax`` and ``auto``.
         transform_program(.TransformProgram): A transform program to be applied to the initial tape.
+        inner_transform (.TransformProgram): A transform program to be applied to the tapes in inner execution, inside the ml interface.
         config (qml.devices.ExecutionConfig): A datastructure describing the parameters needed to fully describe the execution.
         grad_on_execution (bool, str): Whether the gradients should be computed on the execution or not. Only applies
             if the device is queried for the gradient; gradient transform
@@ -584,11 +593,10 @@ def cost_fn(params, x):
             raise ValueError("Using postselect_mode='hw-like' is not supported with jax-jit.")
         config.mcm_config.postselect_mode = "fill-shots"
 
-    if transform_program is None:
-        if isinstance(device, qml.devices.Device):
-            transform_program = device.preprocess(config)[0]
-        else:
-            transform_program = qml.transforms.core.TransformProgram()
+    is_gradient_transform = isinstance(gradient_fn, qml.transforms.core.TransformDispatcher)
+    transform_program, inner_transform = _make_transform_programs(
+        device, config, inner_transform, transform_program, is_gradient_transform
+    )
 
     # If caching is desired but an explicit cache is not provided, use an ``LRUCache``.
     if cache is True:
@@ -614,6 +622,7 @@ def cost_fn(params, x):
         device,
         override_shots,
         cache,
+        inner_transform,
         expand_fn,
         config,
         numpy_only=not device_supports_interface_data,
@@ -751,7 +760,9 @@ def device_execute_and_gradients(internal_tapes, **gradient_kwargs):
 
         else:
             # need to override to have no cache
-            inner_execute = _make_inner_execute(device, override_shots, cache=None)
+            inner_execute = _make_inner_execute(
+                device, override_shots, cache=None, inner_transform=inner_transform
+            )
 
             def inner_execute_with_empty_jac(tapes, **_):
                 return (inner_execute(tapes), [])
@@ -827,6 +838,35 @@ def device_gradient_fn(inner_tapes, **gradient_kwargs):
     return post_processing(results)
 
 
+def _make_transform_programs(
+    device, config, inner_transform, transform_program, is_gradient_transform
+):
+    """helper function to make the transform programs."""
+
+    if isinstance(device, qml.devices.Device):
+
+        # If gradient_fn is a gradient transform, device preprocessing should happen in
+        # inner execute (inside the ml boundary).
+        if is_gradient_transform:
+            if inner_transform is None:
+                inner_transform = device.preprocess(config)[0]
+            if transform_program is None:
+                transform_program = qml.transforms.core.TransformProgram()
+        else:
+            if inner_transform is None:
+                inner_transform = qml.transforms.core.TransformProgram()
+            if transform_program is None:
+                transform_program = device.preprocess(config)[0]
+
+    else:
+        if transform_program is None:
+            transform_program = qml.transforms.core.TransformProgram()
+        if inner_transform is None:
+            inner_transform = qml.transforms.core.TransformProgram()
+
+    return transform_program, inner_transform
+
+
 def _get_execution_config(
     gradient_fn, grad_on_execution, interface, device, device_vjp, mcm_config
 ):