Add adjoint differentiation method

.github/CHANGELOG.md

@@ -2,6 +2,33 @@
 
 <h3>New features since last release</h3>
 
+* A new differentiation method has been added for use with simulators in tape mode. The `"adjoint"`
+  method operates after a forward pass by iteratively applying the inverse gate to scan backwards
+  through the circuit. This method is similar to the reversible method, but has a lower time
+  overhead and a similar memory overhead. It follows the approach provided by
+  [Jones and Gacon](https://arxiv.org/abs/2009.02823).
+
+  Example use:
+
+  ```python
+  import pennylane as qml
+
+  qml.enable_tape()
+
+  wires = 1
+  device = qml.device("default.qubit", wires=wires)
+
+  @qml.qnode(device, diff_method="adjoint")
+  def f(params):
+      qml.RX(0.1, wires=0)
+      qml.Rot(*params, wires=0)
+      qml.RX(-0.3, wires=0)
+      return qml.expval(qml.PauliZ(0))
+
+  params = [0.1, 0.2, 0.3]
+  qml.grad(f)(params)
+  ```
+
 * Added `qml.math.squeeze`.
   [(#1011)](https://github.com/PennyLaneAI/pennylane/pull/1011)
 

pennylane/_qubit_device.py

@@ -24,9 +24,10 @@
 
 import numpy as np
 
+import pennylane as qml
 from pennylane.operation import Sample, Variance, Expectation, Probability, State
 from pennylane.qnodes import QuantumFunctionError
-from pennylane import Device
+from pennylane import Device, math
 from pennylane.wires import Wires
 
 
@@ -688,3 +689,110 @@ def sample(self, observable):
         unraveled_indices = [2] * len(device_wires)
         indices = np.ravel_multi_index(samples.T, unraveled_indices)
         return observable.eigvals[indices]
+
+    def adjoint_jacobian(self, tape):
+        """Implements the method outlined in https://arxiv.org/abs/2009.02823 to calculate the
+        Jacobian."""
+
+        for m in tape.measurements:
+            if m.obs is None:
+                raise ValueError(f"Adjoint differentiation method does not support measurement {m}")
+
+            if not hasattr(m.obs, "base_name"):
+                m.obs.base_name = None  # This is needed for when the observable is a tensor product
+
+        # Perform the forward pass
+        self.reset()
+
+        # Consider using caching and calling lower-level functionality. We just need the state
+        # without postprocessing https://github.com/PennyLaneAI/pennylane/pull/1032/files#r563441040
+        self.execute(tape)
+
+        phi = self._reshape(self.state, [2] * self.num_wires)
+
+        lambdas = [self._apply_operation(phi, obs) for obs in tape.observables]
+
+        jac = np.zeros((len(tape.observables), len(tape.trainable_params)))
+
+        expanded_ops = []
+        for op in reversed(tape.operations):
+            if op.num_params > 1:
+                if isinstance(op, qml.Rot) and not op.inverse:
+                    ops = op.decomposition(*op.parameters, wires=op.wires)
+                    expanded_ops.extend(reversed(ops))
+                else:
+                    raise QuantumFunctionError(
+                        f"The {op.name} operation is not supported using "
+                        'the "adjoint" differentiation method'
+                    )
+            else:
+                expanded_ops.append(op)
+
+        expanded_ops = [o for o in expanded_ops if not o.name in ("QubitStateVector", "BasisState")]
+        dot_product_real = lambda a, b: self._real(math.sum(self._conj(a) * b))
+
+        param_number = len(tape._par_info) - 1
+        trainable_param_number = len(tape.trainable_params) - 1
+        for op in expanded_ops:
+
+            if op.grad_method and param_number in tape.trainable_params:
+                d_op_matrix = operation_derivative(op)
+
+            op.inv()
+            phi = self._apply_operation(phi, op)
+
+            if op.grad_method:
+                if param_number in tape.trainable_params:
+                    mu = self._apply_unitary(phi, d_op_matrix, op.wires)
+
+                    jac_column = np.array(
+                        [2 * dot_product_real(lambda_, mu) for lambda_ in lambdas]
+                    )
+                    jac[:, trainable_param_number] = jac_column
+                    trainable_param_number -= 1
+                param_number -= 1
+
+            lambdas = [self._apply_operation(lambda_, op) for lambda_ in lambdas]
+            op.inv()
+
+        return jac
+
+
+def operation_derivative(operation) -> np.ndarray:
+    r"""Calculate the derivative of an operation.
+
+    For an operation :math:`e^{i \hat{H} \phi t}`, this function returns the matrix representation
+    in the standard basis of its derivative with respect to :math:`t`, i.e.,
+
+    .. math:: \frac{d \, e^{i \hat{H} phi t}}{dt} = i \phi \hat{H} e^{i \hat{H} phi t}.
+
+    Args:
+        operation (qml.Operation): The operation to be differentiated.
+
+    Returns:
+        np.ndarray: the derivative of the operation as a matrix in the standard basis
+
+    Raises:
+        ValueError: if the operation does not have a generator or is not composed of a single
+        trainable parameter
+    """
+    generator, prefactor = operation.generator
+
+    if generator is None:
+        raise ValueError(f"Operation {operation.name} does not have a generator")
+    if operation.num_params != 1:
+        # Note, this case should already be caught by the previous raise since we haven't worked out
+        # how to have an operator for multiple parameters. It is added here in case of a future
+        # change
+        raise ValueError(
+            f"Operation {operation.name} is not written in terms of a single parameter"
+        )
+
+    if not isinstance(generator, np.ndarray):
+        generator = generator.matrix
+
+    if operation.inverse:
+        prefactor *= -1
+        generator = generator.conj().T
+
+    return 1j * prefactor * generator @ operation.matrix

pennylane/tape/qnode.py

@@ -103,6 +103,12 @@ class QNode:
               Only allowed on (simulator) devices with the "reversible" capability,
               for example :class:`default.qubit <~.DefaultQubit>`.
 
+            * ``"adjoint"``: Uses an adjoint `method <https://arxiv.org/abs/2009.02823>`__ that
+              reverses through the circuit after a forward pass by iteratively applying the inverse
+              (adjoint) gate. This method is similar to the reversible method, but has a lower time
+              overhead and a similar memory overhead. Only allowed on simulator devices such as
+              :class:`default.qubit <~.DefaultQubit>`.
+
             * ``"parameter-shift"``: Use the analytic parameter-shift
               rule for all supported quantum operation arguments, with finite-difference
               as a fallback.
@@ -148,15 +154,13 @@ def __init__(self, func, device, interface="autograd", diff_method="best", **dif
         # store the user-specified differentiation method
         self.diff_method = diff_method
 
-        self._tape, self.interface, diff_method, self.device = self.get_tape(
+        self._tape, self.interface, self.device, tape_diff_options = self.get_tape(
             device, interface, diff_method
         )
+
         # The arguments to be passed to JacobianTape.jacobian
         self.diff_options = diff_options or {}
-        # Store the differentiation method to be passed to JacobianTape.jacobian().
-        # Note that the tape accepts a different set of allowed methods than the QNode:
-        #     best, analytic, numeric, device
-        self.diff_options["method"] = diff_method
+        self.diff_options.update(tape_diff_options)
 
         self.dtype = np.float64
         self.max_expansion = 2
@@ -170,7 +174,7 @@ def get_tape(device, interface, diff_method="best"):
             device (.Device): PennyLane device
             interface (str): name of the requested interface
             diff_method (str): The requested method of differentiation. One of
-                ``"best"``, ``"backprop"``, ``"reversible"``, ``"device"``,
+                ``"best"``, ``"backprop"``, ``"reversible"``, ``"adjoint"``, ``"device"``,
                 ``"parameter-shift"``, or ``"finite-diff"``.
 
         Returns:
@@ -188,18 +192,26 @@ def get_tape(device, interface, diff_method="best"):
         if diff_method == "reversible":
             return QNode._validate_reversible_method(device, interface)
 
+        if diff_method == "adjoint":
+            return QNode._validate_adjoint_method(device, interface)
+
         if diff_method == "device":
             return QNode._validate_device_method(device, interface)
 
         if diff_method == "parameter-shift":
-            return QNode._get_parameter_shift_tape(device), interface, "analytic", device
+            return (
+                QNode._get_parameter_shift_tape(device),
+                interface,
+                device,
+                {"method": "analytic"},
+            )
 
         if diff_method == "finite-diff":
-            return JacobianTape, interface, "numeric", device
+            return JacobianTape, interface, device, {"method": "numeric"}
 
         raise qml.QuantumFunctionError(
             f"Differentiation method {diff_method} not recognized. Allowed "
-            "options are ('best', 'parameter-shift', 'backprop', 'finite-diff', 'device', 'reversible')."
+            "options are ('best', 'parameter-shift', 'backprop', 'finite-diff', 'device', 'reversible', 'adjoint')."
         )
 
     @staticmethod
@@ -234,9 +246,14 @@ def get_best_method(device, interface):
                 return QNode._validate_backprop_method(device, interface)
             except qml.QuantumFunctionError:
                 try:
-                    return QNode._get_parameter_shift_tape(device), interface, "best", device
+                    return (
+                        QNode._get_parameter_shift_tape(device),
+                        interface,
+                        device,
+                        {"method": "best"},
+                    )
                 except qml.QuantumFunctionError:
-                    return JacobianTape, interface, "numeric", device
+                    return JacobianTape, interface, device, {"method": "numeric"}
 
     @staticmethod
     def _validate_backprop_method(device, interface):
@@ -271,7 +288,7 @@ def _validate_backprop_method(device, interface):
             # device supports backpropagation natively
 
             if interface == backprop_interface:
-                return JacobianTape, interface, "backprop", device
+                return JacobianTape, interface, device, {"method": "backprop"}
 
             raise qml.QuantumFunctionError(
                 f"Device {device.short_name} only supports diff_method='backprop' when using the "
@@ -284,8 +301,13 @@ def _validate_backprop_method(device, interface):
             if interface in backprop_devices:
                 # TODO: need a better way of passing existing device init options
                 # to a new device?
-                device = qml.device(backprop_devices[interface], wires=device.wires, analytic=True)
-                return JacobianTape, interface, "backprop", device
+                device = qml.device(
+                    backprop_devices[interface],
+                    wires=device.wires,
+                    shots=device.shots,
+                    analytic=True,
+                )
+                return JacobianTape, interface, device, {"method": "backprop"}
 
             raise qml.QuantumFunctionError(
                 f"Device {device.short_name} only supports diff_method='backprop' when using the "
@@ -323,7 +345,41 @@ def _validate_reversible_method(device, interface):
                 f"The {device.short_name} device does not support reversible differentiation."
             )
 
-        return ReversibleTape, interface, "analytic", device
+        return ReversibleTape, interface, device, {"method": "analytic"}
+
+    @staticmethod
+    def _validate_adjoint_method(device, interface):
+        """Validates whether a particular device and JacobianTape interface
+        supports the ``"adjoint"`` differentiation method.
+
+        Args:
+            device (.Device): PennyLane device
+            interface (str): name of the requested interface
+
+        Returns:
+            tuple[.JacobianTape, str, str]: tuple containing the compatible
+            JacobianTape, the interface to apply, and the method argument
+            to pass to the ``JacobianTape.jacobian`` method
+
+        Raises:
+            qml.QuantumFunctionError: if the device does not support adjoint backprop
+        """
+        supported_device = hasattr(device, "_apply_operation")
+        supported_device = supported_device and hasattr(device, "_apply_unitary")
+        supported_device = supported_device and device.capabilities().get("returns_state")
+        supported_device = supported_device and hasattr(device, "adjoint_jacobian")
+
+        if not supported_device:
+            raise ValueError(
+                f"The {device.short_name} device does not support adjoint differentiation."
+            )
+
+        return (
+            JacobianTape,
+            interface,
+            device,
+            {"method": "device", "jacobian_method": "adjoint_jacobian"},
+        )
 
     @staticmethod
     def _validate_device_method(device, interface):
@@ -352,7 +408,7 @@ def _validate_device_method(device, interface):
                 "method for computing the jacobian."
             )
 
-        return JacobianTape, interface, "device", device
+        return JacobianTape, interface, device, {"method": "device"}
 
     @staticmethod
     def _get_parameter_shift_tape(device):
@@ -591,12 +647,12 @@ def to_tf(self, dtype=None):
 
             if self.interface != "tf" and self.interface is not None:
                 # Since the interface is changing, need to re-validate the tape class.
-                self._tape, interface, diff_method, self.device = self.get_tape(
+                self._tape, interface, self.device, diff_options = self.get_tape(
                     self._original_device, "tf", self.diff_method
                 )
 
                 self.interface = interface
-                self.diff_options["method"] = diff_method
+                self.diff_options.update(diff_options)
             else:
                 self.interface = "tf"
 
@@ -631,12 +687,12 @@ def to_torch(self, dtype=None):
 
             if self.interface != "torch" and self.interface is not None:
                 # Since the interface is changing, need to re-validate the tape class.
-                self._tape, interface, diff_method, self.device = self.get_tape(
+                self._tape, interface, self.device, diff_options = self.get_tape(
                     self._original_device, "torch", self.diff_method
                 )
 
                 self.interface = interface
-                self.diff_options["method"] = diff_method
+                self.diff_options.update(diff_options)
             else:
                 self.interface = "torch"
 
@@ -663,12 +719,12 @@ def to_autograd(self):
 
         if self.interface != "autograd" and self.interface is not None:
             # Since the interface is changing, need to re-validate the tape class.
-            self._tape, interface, diff_method, self.device = self.get_tape(
+            self._tape, interface, self.device, diff_options = self.get_tape(
                 self._original_device, "autograd", self.diff_method
             )
 
             self.interface = interface
-            self.diff_options["method"] = diff_method
+            self.diff_options.update(diff_options)
         else:
             self.interface = "autograd"
 
@@ -757,6 +813,12 @@ def qnode(device, interface="autograd", diff_method="best", **diff_options):
               Only allowed on (simulator) devices with the "reversible" capability,
               for example :class:`default.qubit <~.DefaultQubit>`.
 
+            * ``"adjoint"``: Uses an adjoint `method <https://arxiv.org/abs/2009.02823>`__ that
+              reverses through the circuit after a forward pass by iteratively applying the inverse
+              (adjoint) gate. This method is similar to the reversible method, but has a lower time
+              overhead and a similar memory overhead. Only allowed on simulator devices such as
+              :class:`default.qubit <~.DefaultQubit>`.
+
             * ``"device"``: Queries the device directly for the gradient.
               Only allowed on devices that provide their own gradient rules.
 

pennylane/tape/tapes/jacobian_tape.py

@@ -332,7 +332,8 @@ def device_pd(self, device, params=None, **options):
             params (list[Any]): The quantum tape operation parameters. If not provided,
                 the current tape parameter values are used (via :meth:`~.get_parameters`).
         """
-        # pylint:disable=unused-argument
+        jacobian_method = getattr(device, options.get("jacobian_method", "jacobian"))
+
         if params is None:
             params = np.array(self.get_parameters())
 
@@ -343,7 +344,7 @@ def device_pd(self, device, params=None, **options):
 
         # TODO: modify devices that have device Jacobian methods to
         # accept the quantum tape as an argument
-        jac = device.jacobian(self)
+        jac = jacobian_method(self)
 
         # restore original parameters
         self.set_parameters(saved_parameters)