Wrapping of legacy device automatically in various device creation/qnode/execute functions

pennylane/devices/legacy_facade.py

@@ -22,10 +22,9 @@
 from dataclasses import replace
 
 import pennylane as qml
-from pennylane.measurements import Shots
+from pennylane.measurements import MidMeasureMP, Shots
 from pennylane.transforms.core.transform_program import TransformProgram
 
-from .default_qubit import adjoint_observables, adjoint_ops
 from .device_api import Device
 from .execution_config import DefaultExecutionConfig
 from .modifiers import single_tape_support
@@ -34,10 +33,16 @@
     no_sampling,
     validate_adjoint_trainable_params,
     validate_measurements,
-    validate_observables,
 )
 
 
+def _requests_adjoint(execution_config):
+    return execution_config.gradient_method == "adjoint" or (
+        execution_config.gradient_method == "device"
+        and execution_config.gradient_keyword_arguments.get("method", None) == "adjoint_jacobian"
+    )
+
+
 @contextmanager
 def _set_shots(device, shots):
     """Context manager to temporarily change the shots
@@ -98,6 +103,15 @@ def legacy_device_batch_transform(tape, device):
     return _set_shots(device, tape.shots)(device.batch_transform)(tape)
 
 
+def adjoint_ops(op: qml.operation.Operator) -> bool:
+    """Specify whether or not an Operator is supported by adjoint differentiation."""
+    if isinstance(op, qml.QubitUnitary) and not qml.operation.is_trainable(op):
+        return True
+    return not isinstance(op, MidMeasureMP) and (
+        op.num_params == 0 or (op.num_params == 1 and op.has_generator)
+    )
+
+
 def _add_adjoint_transforms(program: TransformProgram, name="adjoint"):
     """Add the adjoint specific transforms to the transform program."""
     program.add_transform(no_sampling, name=name)
@@ -106,7 +120,6 @@ def _add_adjoint_transforms(program: TransformProgram, name="adjoint"):
         stopping_condition=adjoint_ops,
         name=name,
     )
-    program.add_transform(validate_observables, adjoint_observables, name=name)
 
     def accepted_adjoint_measurements(mp):
         return isinstance(mp, qml.measurements.ExpectationMP)
@@ -144,11 +157,11 @@ class LegacyDeviceFacade(Device):
 
     """
 
-    def __new__(cls, device: "qml.devices.LegacyDevice", *args, **kwargs):
-        return device if isinstance(device, cls) else super().__new__(cls)
-
     # pylint: disable=super-init-not-called
     def __init__(self, device: "qml.devices.LegacyDevice"):
+        if isinstance(device, type(self)):
+            raise RuntimeError("An already-facaded device can not be wrapped in a facade again.")
+
         if not isinstance(device, qml.devices.LegacyDevice):
             raise ValueError(
                 "The LegacyDeviceFacade only accepts a device of type qml.devices.LegacyDevice."
@@ -208,11 +221,7 @@ def preprocess(self, execution_config=DefaultExecutionConfig):
         program.add_transform(legacy_device_batch_transform, device=self._device)
         program.add_transform(legacy_device_expand_fn, device=self._device)
 
-        if execution_config.gradient_method == "adjoint" or (
-            execution_config.gradient_method == "device"
-            and execution_config.gradient_keyword_arguments.get("method", None)
-            == "adjoint_jacobian"
-        ):
+        if _requests_adjoint(execution_config):
             _add_adjoint_transforms(program, name=f"{self.name} + adjoint")
 
         if self._device.capabilities().get("supports_mid_measure", False):
@@ -249,9 +258,6 @@ def _setup_adjoint_config(self, execution_config):
         return replace(execution_config, **updated_values)
 
     def _setup_device_config(self, execution_config):
-        if execution_config.gradient_keyword_arguments.get("method", None) == "adjoint_jacobian":
-            return self._setup_adjoint_config(execution_config)
-
         tape = qml.tape.QuantumScript([], [])
 
         if not self._validate_device_method(tape):
@@ -276,13 +282,9 @@ def _setup_execution_config(self, execution_config):
                 config = replace(execution_config, gradient_method="backprop")
                 return self._setup_backprop_config(config)
 
-            if self._validate_adjoint_method(tape):
-                config = replace(execution_config, gradient_method="adjoint")
-                return self._setup_adjoint_config(config)
-
         if execution_config.gradient_method == "backprop":
             return self._setup_backprop_config(execution_config)
-        if execution_config.gradient_method == "adjoint":
+        if _requests_adjoint(execution_config):
             return self._setup_adjoint_config(execution_config)
         if execution_config.gradient_method == "device":
             return self._setup_device_config(execution_config)
@@ -295,14 +297,13 @@ def supports_derivatives(self, execution_config=None, circuit=None) -> bool:
         if execution_config is None or execution_config.gradient_method == "best":
             validation_methods = (
                 self._validate_backprop_method,
-                self._validate_adjoint_method,
                 self._validate_device_method,
             )
             return any(validate(circuit) for validate in validation_methods)
 
         if execution_config.gradient_method == "backprop":
             return self._validate_backprop_method(circuit)
-        if execution_config.gradient_method == "adjoint":
+        if _requests_adjoint(execution_config):
             return self._validate_adjoint_method(circuit)
         if execution_config.gradient_method == "device":
             return self._validate_device_method(circuit)

pennylane/optimize/qnspsa.py

@@ -442,15 +442,11 @@ def _apply_blocking(self, cost, args, kwargs, params_next):
         cost.construct(params_next, kwargs)
         tape_loss_next = cost.tape.copy(copy_operations=True)
 
-        if isinstance(cost.device, qml.devices.Device):
-            program, _ = cost.device.preprocess()
-
-            loss_curr, loss_next = qml.execute(
-                [tape_loss_curr, tape_loss_next], cost.device, None, transform_program=program
-            )
+        program, _ = cost.device.preprocess()
 
-        else:
-            loss_curr, loss_next = qml.execute([tape_loss_curr, tape_loss_next], cost.device, None)
+        loss_curr, loss_next = qml.execute(
+            [tape_loss_curr, tape_loss_next], cost.device, None, transform_program=program
+        )
 
         # self.k has been updated earlier
         ind = (self.k - 2) % self.last_n_steps.size

pennylane/optimize/spsa.py

@@ -265,12 +265,8 @@ def compute_grad(self, objective_fn, args, kwargs):
         try:
             # pylint: disable=protected-access
             dev_shots = objective_fn.device.shots
-            if isinstance(dev_shots, Shots):
-                shots = dev_shots if dev_shots.has_partitioned_shots else Shots(None)
-            elif objective_fn.device.shot_vector is not None:
-                shots = Shots(objective_fn.device._raw_shot_sequence)  # pragma: no cover
-            else:
-                shots = Shots(None)
+
+            shots = dev_shots if dev_shots.has_partitioned_shots else Shots(None)
 
             if np.prod(objective_fn.func(*args, **kwargs).shape(objective_fn.device, shots)) > 1:
                 raise ValueError(

pennylane/transforms/core/transform_dispatcher.py

@@ -306,9 +306,7 @@ def original_device(self):
                 """Return the original device."""
                 return self._original_device
 
-        new_dev = TransformedDevice(original_device, self._transform)
-
-        return new_dev
+        return TransformedDevice(original_device, self._transform)
 
     def _batch_transform(self, original_batch, targs, tkwargs):
         """Apply the transform on a batch of tapes."""

pennylane/workflow/construct_batch.py

@@ -65,14 +65,8 @@ def _get_full_transform_program(qnode: QNode) -> "qml.transforms.core.TransformP
             **qnode.gradient_kwargs,
         )
 
-    if isinstance(qnode.device, qml.devices.Device):
-        config = _make_execution_config(qnode, qnode.gradient_fn)
-        return program + qnode.device.preprocess(config)[0]
-
-    program.add_transform(qml.transform(qnode.device.batch_transform))
-    program.add_transform(expand_fn_transform(qnode.device.expand_fn))
-
-    return program
+    config = _make_execution_config(qnode, qnode.gradient_fn)
+    return program + qnode.device.preprocess(config)[0]
 
 
 def get_transform_program(qnode: "QNode", level=None) -> "qml.transforms.core.TransformProgram":

pennylane/workflow/execution.py

@@ -663,7 +663,7 @@ def inner_execute_with_empty_jac(tapes, **_):
         device_vjp
         and getattr(device, "short_name", "") in ("lightning.gpu", "lightning.kokkos")
         and interface in jpc_interfaces
-    ):
+    ):  # pragma: no cover
         if INTERFACE_MAP[interface] == "jax" and "use_device_state" in gradient_kwargs:
             gradient_kwargs["use_device_state"] = False
 
@@ -781,26 +781,18 @@ def _make_transform_programs(
 ):
     """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 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]
 
     return transform_program, inner_transform
 

pennylane/workflow/qnode.py

@@ -20,10 +20,13 @@
 import inspect
 import logging
 import warnings
-from collections.abc import Sequence
-from typing import Optional
+from collections.abc import Callable, Sequence
+from typing import Any, Literal, Optional, Union, get_args
+
+from cachetools import Cache
 
 import pennylane as qml
+from pennylane import Device
 from pennylane.debugging import pldb_device_manager
 from pennylane.logging import debug_logger
 from pennylane.measurements import CountsMP, MidMeasureMP
@@ -47,6 +50,8 @@
     "spsa",
 ]
 
+SupportedDeviceAPIs = Union[Device, "qml.devices.Device"]
+
 
 def _convert_to_interface(res, interface):
     """
@@ -454,19 +459,19 @@ def circuit_unpacking(x):
 
     def __init__(
         self,
-        func,
-        device: "qml.devices.Device",
-        interface="auto",
-        diff_method="best",
-        expansion_strategy=None,
-        max_expansion=None,
-        grad_on_execution="best",
-        cache="auto",
-        cachesize=10000,
-        max_diff=1,
-        device_vjp=False,
-        postselect_mode=None,
-        mcm_method=None,
+        func: Callable,
+        device: SupportedDeviceAPIs,
+        interface: SupportedInterfaceUserInput = "auto",
+        diff_method: Union[TransformDispatcher, SupportedDiffMethods] = "best",
+        expansion_strategy: Literal[None, "device", "gradient"] = None,
+        max_expansion: Optional[int] = None,
+        grad_on_execution: Literal[True, False, "best"] = "best",
+        cache: Union[Cache, Literal["auto", True, False]] = "auto",
+        cachesize: int = 10000,
+        max_diff: int = 1,
+        device_vjp: Union[None, bool] = False,
+        postselect_mode: Literal[None, "hw-like", "fill-shots"] = None,
+        mcm_method: Literal[None, "deferred", "one-shot", "tree-traversal"] = None,
         **gradient_kwargs,
     ):
         # Moving it here since the old default value is checked on debugging
@@ -508,9 +513,6 @@ def __init__(
                 gradient_kwargs,
             )
 
-        if not isinstance(device, qml.devices.Device):
-            device = qml.devices.LegacyDeviceFacade(device)
-
         if interface not in SUPPORTED_INTERFACES:
             raise qml.QuantumFunctionError(
                 f"Unknown interface {interface}. Interface must be "
@@ -522,6 +524,9 @@ def __init__(
                 "Invalid device. Device must be a valid PennyLane device."
             )
 
+        if not isinstance(device, qml.devices.Device):
+            device = qml.devices.LegacyDeviceFacade(device)
+
         if "shots" in inspect.signature(func).parameters:
             warnings.warn(
                 "Detected 'shots' as an argument to the given quantum function. "
@@ -668,7 +673,7 @@ def _update_gradient_fn(self, shots=None, tape: Optional["qml.tape.QuantumTape"]
     @staticmethod
     @debug_logger
     def get_gradient_fn(
-        device: Union[Device, "qml.devices.Device"],
+        device: SupportedDeviceAPIs,
         interface,
         diff_method: Union[TransformDispatcher, SupportedDiffMethods] = "best",
         tape: Optional["qml.tape.QuantumTape"] = None,
@@ -734,13 +739,13 @@ def get_gradient_fn(
     @staticmethod
     @debug_logger
     def get_best_method(
-        device: Union[Device, "qml.devices.Device"],
-        interface,
+        device: SupportedDeviceAPIs,
+        interface: SupportedInterfaceUserInput,
         tape: Optional["qml.tape.QuantumTape"] = None,
     ) -> tuple[
         Union[TransformDispatcher, Literal["device", "backprop", "parameter-shift", "finite-diff"]],
         dict[str, Any],
-        Union[Device, "qml.devices.Device"],
+        SupportedDeviceAPIs,
     ]:
         """Returns the 'best' differentiation method
         for a particular device and interface combination.
@@ -782,7 +787,7 @@ def get_best_method(
 
     @staticmethod
     @debug_logger
-    def best_method_str(device: Union[Device, "qml.devices.Device"], interface) -> str:
+    def best_method_str(device: SupportedDeviceAPIs, interface: SupportedInterfaceUserInput) -> str:
         """Similar to :meth:`~.get_best_method`, except return the
         'best' differentiation method in human-readable format.
 

tests/devices/test_legacy_facade.py

@@ -53,8 +53,16 @@ def expval(self, observable, wires, par):
         return 0.0
 
 
+def test_double_facade_raises_error():
+    """Test that a RuntimeError is raised if a facaded device is passed to constructor"""
+    dev = qml.device("default.mixed", wires=1)
+
+    with pytest.raises(RuntimeError, match="already-facaded device can not be wrapped"):
+        qml.devices.LegacyDeviceFacade(dev)
+
+
 def test_error_if_not_legacy_device():
-    """Test that a ValueError is raiuised if the target is not a legacy device."""
+    """Test that a ValueError is raised if the target is not a legacy device."""
 
     target = qml.devices.DefaultQubit()
     with pytest.raises(ValueError, match="The LegacyDeviceFacade only accepts"):
@@ -282,10 +290,11 @@ def test_no_derivatives_case(self):
         with pytest.raises(qml.DeviceError):
             dev.preprocess(ExecutionConfig(gradient_method="backprop"))
 
-    @pytest.mark.parametrize("gradient_method", ("best", "adjoint"))
-    def test_adjoint_support(self, gradient_method):
+    def test_adjoint_support(self):
         """Test that the facade can handle devices that support adjoint."""
 
+        gradient_method = "adjoint"
+
         # pylint: disable=unnecessary-lambda-assignment
         class AdjointDev(DummyDevice):
             """A dummy device that supports adjoint diff"""

tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_legacy.py

@@ -157,10 +157,10 @@ def circuit(p1, p2=y, **kwargs):
             qml.RY(p2[0] * p2[1], wires=1)
             qml.RX(kwargs["p3"], wires=0)
             qml.CNOT(wires=[0, 1])
-            return qml.state()
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
 
         result = qml.draw(circuit)(p1=x, p3=z)
-        expected = "0: ──RX(0.10)──RX(0.40)─╭●─┤  State\n1: ──RY(0.06)───────────╰X─┤  State"
+        expected = "0: ──RX(0.10)──RX(0.40)─╭●─┤  <Z>\n1: ──RY(0.06)───────────╰X─┤  <Z>"
         assert result == expected
 
     def test_jacobian(self, dev_name, diff_method, grad_on_execution, tol, interface):