Adjoint metric tensor (#1992)

* main commit and tests

* actually add files. changelog.

* changelog

* undo scatter_element_add modification as it was not successful

* support new device for qnode wrapper

* black

* linting

* warning test

* black

* lint

* test coverage inverse

* ignore linting issues from other PRs

* Apply suggestions from code review

Co-authored-by: Josh Izaac <josh146@gmail.com>

* code review

* separate cjac contraction from QNode wrapper

* Update changelog-dev.md

* Apply suggestions from code review

Co-authored-by: Josh Izaac <josh146@gmail.com>

* code review

* tests

* black

* tests, output shapes of metric_tensor

* more test adaptation

* test fixes

* rename back to adjoint

* remove import

* remove prints

* Apply suggestions from code review

Co-authored-by: Josh Izaac <josh146@gmail.com>

* code review

* breaking change+example

* revert breaking change in metric_tensor output shape

* black

* reduce code length

* remove "breaking change" comment

* undo some commenting

* Update doc/releases/changelog-dev.md

Co-authored-by: Josh Izaac <josh146@gmail.com>

Co-authored-by: Josh Izaac <josh146@gmail.com>

doc/releases/changelog-dev.md

@@ -4,6 +4,60 @@
 
 <h3>New features since last release</h3>
 
+* Added the adjoint method for the metric tensor.
+  [(#1992)](https://github.com/PennyLaneAI/pennylane/pull/1992)
+
+  This method, detailed in [Jones 2020](https://arxiv.org/abs/2011.02991),
+  computes the metric tensor using four copies of the state vector and
+  a number of operations that scales quadratically in the number of trainable
+  parameters (see below for details).
+
+  Note that as it makes use of state cloning, it is inherently classical
+  and can only be used with statevector simulators and `shots=None`.
+
+  It is particular useful for larger circuits for which backpropagation requires
+  inconvenient or even unfeasible amounts of storage, but is slower.
+  Furthermore, the adjoint method is only available for analytic computation, not
+  for measurements simulation with `shots!=None`.
+
+  ```python
+  dev = qml.device("default.qubit", wires=3)
+
+  @qml.qnode(dev)
+  def circuit(x, y):
+      qml.Rot(*x[0], wires=0)
+      qml.Rot(*x[1], wires=1)
+      qml.Rot(*x[2], wires=2)
+      qml.CNOT(wires=[0, 1])
+      qml.CNOT(wires=[1, 2])
+      qml.CNOT(wires=[2, 0])
+      qml.RY(y[0], wires=0)
+      qml.RY(y[1], wires=1)
+      qml.RY(y[0], wires=2)
+
+  x = np.array([[0.2, 0.4, -0.1], [-2.1, 0.5, -0.2], [0.1, 0.7, -0.6]], requires_grad=False)
+  y = np.array([1.3, 0.2], requires_grad=True)
+  ```
+
+  ```pycon
+  >>> qml.adjoint_metric_tensor(circuit)(x, y)
+  tensor([[ 0.25495723, -0.07086695],
+          [-0.07086695,  0.24945606]], requires_grad=True)
+  ```
+
+  Computational cost
+
+  The adjoint method uses :math:`2P^2+4P+1` gates and state cloning operations if the circuit
+  is composed only of trainable gates, where :math:`P` is the number of trainable operations.
+  If non-trainable gates are included, each of them is applied about :math:`n^2-n` times, where
+  :math:`n` is the number of trainable operations that follow after the respective 
+  non-trainable operation in the circuit. This means that non-trainable gates later in the 
+  circuit are executed less often, making the adjoint method a bit cheaper if such gates
+  appear later.
+  The adjoint method requires memory for 4 independent state vectors, which corresponds roughly
+  to storing a state vector of a system with 2 additional qubits.
+
+
 <h3>Improvements</h3>
 
 * Interferometer is now a class with `shape` method.

pennylane/__init__.py

@@ -53,6 +53,7 @@
 import pennylane.qnode_old
 from pennylane.transforms import (
     adjoint,
+    adjoint_metric_tensor,
     batch_params,
     batch_transform,
     draw,

pennylane/optimize/qng.py

@@ -204,8 +204,15 @@ def step_and_cost(
 
                 metric_tensor_fn = qml.metric_tensor(qnode, approx=self.approx)
 
-            self.metric_tensor = metric_tensor_fn(*args, **kwargs)
-            self.metric_tensor += self.lam * np.identity(self.metric_tensor.shape[0])
+            _metric_tensor = metric_tensor_fn(*args, **kwargs)
+            # Reshape metric tensor to be square
+            shape = qml.math.shape(_metric_tensor)
+            size = qml.math.prod(shape[: len(shape) // 2])
+            self.metric_tensor = qml.math.reshape(_metric_tensor, (size, size))
+            # Add regularization
+            self.metric_tensor = self.metric_tensor + self.lam * qml.math.eye(
+                size, like=_metric_tensor
+            )
 
         g, forward = self.compute_grad(qnode, args, kwargs, grad_fn=grad_fn)
         new_args = np.array(self.apply_grad(g, args), requires_grad=True)

pennylane/templates/subroutines/qpe.py

@@ -146,7 +146,7 @@ def expand(self):
 
         return tape
 
-    def adjoint(self):
+    def adjoint(self):  # pylint: disable=arguments-differ
         adjoint_op = QuantumPhaseEstimation(
             *self.parameters, target_wires=self.target_wires, estimation_wires=self.estimation_wires
         )

pennylane/transforms/__init__.py

@@ -35,6 +35,7 @@
     ~draw_mpl
     ~transforms.get_unitary_matrix
     ~metric_tensor
+    ~adjoint_metric_tensor
     ~specs
     ~transforms.mitigate_with_zne
 
@@ -110,6 +111,7 @@
     ~transforms.create_decomp_expand_fn
     ~transforms.expand_invalid_trainable
     ~transforms.expand_multipar
+    ~transforms.expand_trainable_multipar
     ~transforms.expand_nonunitary_gen
 """
 # Import the decorators first to prevent circular imports when used in other transforms
@@ -125,6 +127,7 @@
 from .hamiltonian_expand import hamiltonian_expand
 from .measurement_grouping import measurement_grouping
 from .metric_tensor import metric_tensor
+from .adjoint_metric_tensor import adjoint_metric_tensor
 from .insert_ops import insert
 from .mitigate import mitigate_with_zne
 from .optimization import (
@@ -144,6 +147,7 @@
     expand_invalid_trainable,
     expand_multipar,
     expand_nonunitary_gen,
+    expand_trainable_multipar,
     create_expand_fn,
     create_decomp_expand_fn,
     set_decomposition,