Add adjoint differentiation method

.github/CHANGELOG.md

@@ -2,6 +2,34 @@
 
 <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 inverse gates 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). This method is only compatible with certain
+  statevector-based devices such as `default.qubit`.
+
+  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,18 @@
 
 import numpy as np
 
-from pennylane.operation import Sample, Variance, Expectation, Probability, State
+import pennylane as qml
+from pennylane.operation import (
+    Sample,
+    Variance,
+    Expectation,
+    Probability,
+    State,
+    operation_derivative,
+)
 from pennylane.qnodes import QuantumFunctionError
 from pennylane import Device
+from pennylane.math import sum
 from pennylane.wires import Wires
 
 
@@ -688,3 +697,96 @@ 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 adjoint method outlined in
+        `Jones and Gacon <https://arxiv.org/abs/2009.02823>`__ to differentiate an input tape.
+
+        After a forward pass, the circuit is reversed by iteratively applying inverse (adjoint)
+        gates 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.
+
+        .. note::
+            The adjoint differentation method has the following restrictions:
+
+            * As it requires knowledge of the statevector, only statevector simulator devices can be
+              used.
+
+            * Only expectation values are supported as measurements.
+
+        Args:
+            tape (.QuantumTape): circuit that the function takes the gradient of
+
+        Returns:
+            array: the derivative of the tape with respect to trainable parameters.
+            Dimensions are ``(len(observables), len(trainable_params))``.
+
+        Raises:
+            QuantumFunctionError: if the input tape has measurements that are not expectation values
+                or contains a multi-parameter operation aside from :class:`~.Rot`
+        """
+
+        for m in tape.measurements:
+            if m.return_type is not qml.operation.Expectation:
+                raise qml.QuantumFunctionError(
+                    "Adjoint differentiation method does not support"
+                    f" measurement {m.return_type.value}"
+                )
+
+            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]
+
+        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:
+                if op.name not in ("QubitStateVector", "BasisState"):
+                    expanded_ops.append(op)
+
+        jac = np.zeros((len(tape.observables), len(tape.trainable_params)))
+        dot_product_real = lambda a, b: self._real(sum(self._conj(a) * b))
+
+        param_number = len(tape._par_info) - 1  # pylint: disable=protected-access
+        trainable_param_number = len(tape.trainable_params) - 1
+        for op in expanded_ops:
+
+            if (op.grad_method is not None) 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 is not None:
+                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

pennylane/operation.py

@@ -1727,3 +1727,45 @@ def heisenberg_obs(self, wires):
         p = self.parameters
         U = self._heisenberg_rep(p)  # pylint: disable=assignment-from-none
         return self.heisenberg_expand(U, wires)
+
+
+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},
+
+    where :math:`\phi` is a real constant.
+
+    Args:
+        operation (qml.Operation): The operation to be differentiated.
+
+    Returns:
+        array: 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