Integrate experimental device with the QNode (#4196)

* integrate qnode with new device * some diff method improvements * repr methods * add tests, always pass config to device * add tests, always pass config to device * final test * pylint * autograd integration tests * pylint * Update pennylane/interfaces/execution.py Co-authored-by: Matthew Silverman <matthews@xanadu.ai> * pass shots through methods * changelog * revert set shots change * Apply suggestions from code review Co-authored-by: Matthew Silverman <matthews@xanadu.ai> * revert executionc hange, pylint: * pylint again --------- Co-authored-by: Matthew Silverman <matthews@xanadu.ai> Co-authored-by: Romain Moyard <rmoyard@gmail.com>
PennyLaneAI · Jun 19, 2023 · 0785aec · 0785aec
1 parent 723cdc8
commit 0785aec
Show file tree

Hide file tree

Showing 5 changed files with 2,437 additions and 145 deletions.
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
@@ -64,6 +64,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* The experimental device interface is integrated with the `QNode`.
+  [(#4196)](https://github.com/PennyLaneAI/pennylane/pull/4196)
+
 * `Projector` now accepts a state vector representation, which enables the creation of projectors
   in any basis.
   [(#4192)](https://github.com/PennyLaneAI/pennylane/pull/4192)

diff --git a/pennylane/interfaces/execution.py b/pennylane/interfaces/execution.py
@@ -513,7 +513,16 @@ def cost_fn(params, x):
 
     # the default execution function is batch_execute
     # use qml.interfaces so that mocker can spy on it during testing
-    execute_fn = qml.interfaces.cache_execute(batch_execute, cache, expand_fn=expand_fn)
+    if new_device_interface:
+
+        def device_execution_with_config(tapes):
+            return device.execute(tapes, execution_config=config)
+
+        execute_fn = qml.interfaces.cache_execute(
+            device_execution_with_config, cache, expand_fn=expand_fn
+        )
+    else:
+        execute_fn = qml.interfaces.cache_execute(batch_execute, cache, expand_fn=expand_fn)
 
     _grad_on_execution = False
 

diff --git a/pennylane/qnode.py b/pennylane/qnode.py
@@ -19,6 +19,8 @@
 import inspect
 import warnings
 from collections.abc import Sequence
+from typing import Union
+
 
 import autograd
 
@@ -378,7 +380,7 @@ def circuit_unpacking(x):
     def __init__(
         self,
         func,
-        device,
+        device: Union[Device, "qml.devices.experimental.Device"],
         interface="auto",
         diff_method="best",
         expansion_strategy="gradient",
@@ -396,7 +398,7 @@ def __init__(
                 f"one of {SUPPORTED_INTERFACES}."
             )
 
-        if not isinstance(device, Device):
+        if not isinstance(device, (Device, qml.devices.experimental.Device)):
             raise qml.QuantumFunctionError(
                 "Invalid device. Device must be a valid PennyLane device."
             )
@@ -461,6 +463,9 @@ def __init__(
 
     def __repr__(self):
         """String representation."""
+        if isinstance(self.device, qml.devices.experimental.Device):
+            return f"<QNode: device='{self.device}', interface='{self.interface}', diff_method='{self.diff_method}'>"
+
         detail = "<QNode: wires={}, device='{}', interface='{}', diff_method='{}'>"
         return detail.format(
             self.device.num_wires,
@@ -499,7 +504,7 @@ def add_transform(self, transform_container):
         """
         self._transform_program.push_back(transform_container=transform_container)
 
-    def _update_gradient_fn(self):
+    def _update_gradient_fn(self, shots=None):
         if self.diff_method is None:
             self._interface = None
             self.gradient_fn = None
@@ -508,16 +513,17 @@ def _update_gradient_fn(self):
         if self.interface == "auto" and self.diff_method in ["backprop", "best"]:
             if self.diff_method == "backprop":
                 # Check that the device has the capabilities to support backprop
-                backprop_devices = self.device.capabilities().get("passthru_devices", None)
-                if backprop_devices is None:
-                    raise qml.QuantumFunctionError(
-                        f"The {self.device.short_name} device does not support native computations with "
-                        "autodifferentiation frameworks."
-                    )
+                if isinstance(self.device, Device):
+                    backprop_devices = self.device.capabilities().get("passthru_devices", None)
+                    if backprop_devices is None:
+                        raise qml.QuantumFunctionError(
+                            f"The {self.device.short_name} device does not support native computations with "
+                            "autodifferentiation frameworks."
+                        )
             return
 
         self.gradient_fn, self.gradient_kwargs, self.device = self.get_gradient_fn(
-            self._original_device, self.interface, self.diff_method
+            self._original_device, self.interface, self.diff_method, shots=shots
         )
         self.gradient_kwargs.update(self._user_gradient_kwargs or {})
 
@@ -541,7 +547,7 @@ def _update_original_device(self):
 
     # pylint: disable=too-many-return-statements
     @staticmethod
-    def get_gradient_fn(device, interface, diff_method="best"):
+    def get_gradient_fn(device, interface, diff_method="best", shots=None):
         """Determine the best differentiation method, interface, and device
         for a requested device, interface, and diff method.
 
@@ -558,10 +564,10 @@ def get_gradient_fn(device, interface, diff_method="best"):
             ``gradient_kwargs``, and the device to use when calling the execute function.
         """
         if diff_method == "best":
-            return QNode.get_best_method(device, interface)
+            return QNode.get_best_method(device, interface, shots=shots)
 
         if diff_method == "backprop":
-            return QNode._validate_backprop_method(device, interface)
+            return QNode._validate_backprop_method(device, interface, shots=shots)
 
         if diff_method == "adjoint":
             return QNode._validate_adjoint_method(device)
@@ -596,7 +602,7 @@ def get_gradient_fn(device, interface, diff_method="best"):
         )
 
     @staticmethod
-    def get_best_method(device, interface):
+    def get_best_method(device, interface, shots=None):
         """Returns the 'best' differentiation method
         for a particular device and interface combination.
 
@@ -624,7 +630,7 @@ def get_best_method(device, interface):
             return QNode._validate_device_method(device)
         except qml.QuantumFunctionError:
             try:
-                return QNode._validate_backprop_method(device, interface)
+                return QNode._validate_backprop_method(device, interface, shots=shots)
             except qml.QuantumFunctionError:
                 try:
                     return QNode._validate_parameter_shift(device)
@@ -670,10 +676,18 @@ def best_method_str(device, interface):
         return transform
 
     @staticmethod
-    def _validate_backprop_method(device, interface):
-        if device.shots is not None:
+    def _validate_backprop_method(device, interface, shots=None):
+        if shots is not None or getattr(device, "shots", None) is not None:
             raise qml.QuantumFunctionError("Backpropagation is only supported when shots=None.")
 
+        if isinstance(device, qml.devices.experimental.Device):
+            config = qml.devices.experimental.ExecutionConfig(
+                gradient_method="backprop", interface=interface
+            )
+            if device.supports_derivatives(config):
+                return "backprop", {}, device
+            raise qml.QuantumFunctionError(f"Device {device.name} does not support backprop")
+
         mapped_interface = INTERFACE_MAP.get(interface, interface)
 
         # determine if the device supports backpropagation
@@ -731,6 +745,11 @@ def _validate_adjoint_method(device):
         # need to inspect the circuit measurements to ensure only expectation values are taken. This
         # cannot be done here since we don't yet know the composition of the circuit.
 
+        if isinstance(device, qml.devices.experimental.Device):
+            config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+            if device.supports_derivatives(config):
+                return "adjoint", {}, device
+            raise ValueError(f"The {device} device does not support adjoint differentiation.")
         required_attrs = ["_apply_operation", "_apply_unitary", "adjoint_jacobian"]
         supported_device = all(hasattr(device, attr) for attr in required_attrs)
         supported_device = supported_device and device.capabilities().get("returns_state")
@@ -750,17 +769,25 @@ def _validate_adjoint_method(device):
 
     @staticmethod
     def _validate_device_method(device):
-        # determine if the device provides its own jacobian method
-        if device.capabilities().get("provides_jacobian", False):
-            return "device", {}, device
+        if isinstance(device, Device):
+            # determine if the device provides its own jacobian method
+            if device.capabilities().get("provides_jacobian", False):
+                return "device", {}, device
+            name = device.short_name
+        else:
+            config = qml.devices.experimental.ExecutionConfig(gradient_method="device")
+            if device.supports_derivatives(config):
+                return "device", {}, device
+            name = device.name
 
         raise qml.QuantumFunctionError(
-            f"The {device.short_name} device does not provide a native "
-            "method for computing the jacobian."
+            f"The {name} device does not provide a native " "method for computing the jacobian."
         )
 
     @staticmethod
     def _validate_parameter_shift(device):
+        if isinstance(device, qml.devices.experimental.Device):
+            return qml.gradients.param_shift, {}, device
         model = device.capabilities().get("model", None)
 
         if model in {"qubit", "qutrit"}:
@@ -890,14 +917,20 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
 
                 # pylint: disable=not-callable
                 # update the gradient function
-                set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
+                if isinstance(self._original_device, Device):
+                    set_shots(self._original_device, override_shots)(self._update_gradient_fn)()
+                else:
+                    self._update_gradient_fn(shots=override_shots)
 
             else:
-                kwargs["shots"] = (
-                    self._original_device._raw_shot_sequence
-                    if self._original_device._shot_vector
-                    else self._original_device.shots
-                )
+                if isinstance(self._original_device, Device):
+                    kwargs["shots"] = (
+                        self._original_device._raw_shot_sequence
+                        if self._original_device._shot_vector
+                        else self._original_device.shots
+                    )
+                else:
+                    kwargs["shots"] = None
 
         # construct the tape
         self.construct(args, kwargs)