Repository: PennyLaneAI/pennylane
Pull request #1992 (2021-12-06 18:44:08)

Adjoint metric tensor
**Context:**
This implements the adjoint method for the metric tensor as detailed in [Jones 2020](https://arxiv.org/abs/2011.02991)
As the device is used on a low level, it is not implemented as a `tape_transform`. It acts as function with signature `tape->array` on tapes and `QNode->callable[QNode args->array]` on QNodes, so that the usage of `adjoint_metric_tensor` and `metric_tensor` is the same on QNodes:
```pycon
>>> qml.metric_tensor(circuit)(x, y, z)
...
>>> qml.adjoint_metric_tensor(circuit)(x, y, z)
...
```

**Description of the Change:**
Addition of `qml.adjoint_metric_tensor`. It should be discussed whether the file location in `pennylane/transforms/` is appropriate and whether we want any form of alias in `metric_tensor` to link back to `adjoint_metric_tensor`. For the latter, the technical discrepancies seem to prevent such a union of the two functionalities.

Also adds the tape expansion function `qml.transforms.expand_trainable_multipar` 

**Benefits:** New method to compute the metric tensor, in particular useful when on a simulator and backpropagation requires too much memory.

**Possible Drawbacks:** n/a

**Related GitHub Issues:**
There was an old PR of mine: #1387 


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1992](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (b839d44) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/53b3b2b8ca9baff88d4aa3b5b72cb7cf9b8b713b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (53b3b2b) will **decrease** coverage by `1.93%`.
> The diff coverage is `94.92%`.

> :exclamation: Current head b839d44 differs from pull request most recent head b106e0c. Consider uploading reports for the commit b106e0c to get more accurate results
[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1992      +/-   ##
==========================================
- Coverage   98.81%   96.88%   -1.94%     
==========================================
  Files         227      228       +1     
  Lines       17291    17436     +145     
==========================================
- Hits        17086    16892     -194     
- Misses        205      544     +339     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19faW5pdF9fLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/math/single\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvc2luZ2xlX2Rpc3BhdGNoLnB5) | `88.49% <66.66%> (-10.31%)` | :arrow_down: |
| [pennylane/transforms/adjoint\_metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYWRqb2ludF9tZXRyaWNfdGVuc29yLnB5) | `95.45% <95.45%> (ø)` | |
| [pennylane/transforms/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/tape\_expand.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvdGFwZV9leHBhbmQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-95.00%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `80.00% <0.00%> (-20.00%)` | :arrow_down: |
| ... and [21 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [53b3b2b...b106e0c](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1992?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: [sc-12389]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 3f4fd8e8c..31a9314cd 100644
--- a/doc/releases/changelog-dev.md
+++ b/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.
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index b38cc86a8..c30435f98 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -53,6 +53,7 @@ from pennylane.qnode import QNode, qnode
 import pennylane.qnode_old
 from pennylane.transforms import (
     adjoint,
+    adjoint_metric_tensor,
     batch_params,
     batch_transform,
     draw,
diff --git a/pennylane/optimize/qng.py b/pennylane/optimize/qng.py
index 6513b8872..cba9c047b 100644
--- a/pennylane/optimize/qng.py
+++ b/pennylane/optimize/qng.py
@@ -204,8 +204,15 @@ class QNGOptimizer(GradientDescentOptimizer):
 
                 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)
diff --git a/pennylane/templates/subroutines/qpe.py b/pennylane/templates/subroutines/qpe.py
index 9ecc588c0..bba6657ef 100644
--- a/pennylane/templates/subroutines/qpe.py
+++ b/pennylane/templates/subroutines/qpe.py
@@ -146,7 +146,7 @@ class QuantumPhaseEstimation(Operation):
 
         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
         )
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 5b810f613..3b7817589 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -35,6 +35,7 @@ that compute the desired quantity.
     ~draw_mpl
     ~transforms.get_unitary_matrix
     ~metric_tensor
+    ~adjoint_metric_tensor
     ~specs
     ~transforms.mitigate_with_zne
 
@@ -110,6 +111,7 @@ to help build custom QNode, quantum function, and tape transforms:
     ~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 .draw import draw, draw_mpl
 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 @@ from .tape_expand import (
     expand_invalid_trainable,
     expand_multipar,
     expand_nonunitary_gen,
+    expand_trainable_multipar,
     create_expand_fn,
     create_decomp_expand_fn,
     set_decomposition,
diff --git a/pennylane/transforms/adjoint_metric_tensor.py b/pennylane/transforms/adjoint_metric_tensor.py
new file mode 100644
index 000000000..8dfbc22ee
--- /dev/null
+++ b/pennylane/transforms/adjoint_metric_tensor.py
@@ -0,0 +1,308 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Contains the adjoint_metric_tensor.
+"""
+import warnings
+from itertools import chain
+from pennylane import numpy as np
+
+import pennylane as qml
+
+# pylint: disable=protected-access
+from pennylane.transforms.metric_tensor import _contract_metric_tensor_with_cjac
+
+
+def _get_generator(op):
+    """Reads out the generator and prefactor of an operation and converts
+    to matrix if necessary.
+
+    Args:
+        op (:class:`~.Operation`): Operation to obtain the generator of.
+    Returns:
+        array[float]: Generator matrix
+        float: Prefactor of the generator
+    """
+
+    generator, prefactor = op.generator
+    if not isinstance(generator, np.ndarray):
+        generator = generator.matrix
+    if op.inverse:
+        generator = generator.conj().T
+        prefactor *= -1
+
+    return generator, prefactor
+
+
+def _apply_operations(state, op, device, invert=False):
+    """Wrapper that allows to apply a variety of operations---or groups
+    of operations---to a state or to prepare a new state.
+    If ``invert=True``, this function makes sure not to alter the operations.
+    The state of the device, however may be altered, depending on the
+    device and performed operation(s).
+    """
+    # pylint: disable=protected-access
+    if isinstance(op, (list, np.ndarray)):
+        if invert:
+            op = op[::-1]
+        for _op in op:
+            state = _apply_operations(state, _op, device, invert)
+        return state
+
+    if isinstance(op, qml.QubitStateVector):
+        if invert:
+            raise ValueError("Can't invert state preparation.")
+        device._apply_state_vector(op.parameters[0], op.wires)
+        return device._state
+
+    if isinstance(op, qml.BasisState):
+        if invert:
+            raise ValueError("Can't invert state preparation.")
+        device._apply_basis_state(op.parameters[0], op.wires)
+        return device._state
+
+    if invert:
+        op.inv()
+    state = device._apply_operation(state, op)
+    if invert:
+        op.inv()
+
+    return state
+
+
+def _group_operations(tape):
+    """Divide all operations of a tape into trainable operations and blocks
+    of untrainable operations after each trainable one."""
+
+    # Extract tape operations list
+    ops = tape.operations
+    # Find the indices of trainable operations in the tape operations list
+    trainables = np.where([qml.operation.is_trainable(op) for op in ops])[0]
+    # Add the indices incremented by one to the trainable indices
+    split_ids = list(chain.from_iterable([idx, idx + 1] for idx in trainables))
+
+    # Split at trainable and incremented indices to get groups after trainable
+    # operations and single trainable operations (in alternating order)
+    all_groups = np.split(ops, split_ids)
+
+    # Collect trainable operations and groups after trainable operations
+    # the first set of non-trainable ops are the ops "after the -1st" trainable op
+    group_after_trainable_op = dict(enumerate(all_groups[::2], start=-1))
+    trainable_operations = list(chain.from_iterable(all_groups[1::2]))
+
+    return trainable_operations, group_after_trainable_op
+
+
+def adjoint_metric_tensor(circuit, device=None, hybrid=True):
+    r"""Implements the adjoint method outlined in
+    `Jones <https://arxiv.org/abs/2011.02991>`__ to compute the metric tensor.
+
+    A forward pass followed by intermediate partial backwards passes are
+    used to evaluate the metric tensor in :math:`\mathcal{O}(p^2)` operations,
+    where :math:`p` is the number of trainable operations, using 4 state
+    vectors.
+
+    .. note::
+        The adjoint metric tensor method has the following restrictions:
+
+        * Currently only ``"default.qubit"`` with ``shots=None`` is supported.
+
+        * We assume the circuit to be composed of unitary gates only and rely
+          on the ``generator`` property of the gates to be implemented.
+          Note also that this makes the metric tensor strictly real-valued.
+
+    Args:
+        circuit (.QuantumTape or .QNode): Circuit to compute the metric tensor of
+        device (.Device): Device to use for the adjoint method
+        hybrid (bool): Whether to take classical preprocessing into account. Ignored if
+            ``circuit`` is a tape.
+
+    Returns:
+        array: the metric tensor of the tape with respect to its trainable parameters.
+        Dimensions are ``(tape.num_params, tape.num_params)``.
+
+    .. seealso:: :func:`~.metric_tensor` for hardware-compatible metric tensor computations.
+
+    **Example**
+
+    Consider the following QNode:
+
+    .. code-block:: python
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)), qml.expval(qml.PauliY(1))
+
+    We can use the ``adjoint_metric_tensor`` transform to generate a new function
+    that returns the metric tensor of this QNode:
+
+    >>> mt_fn = qml.adjoint_metric_tensor(circuit)
+    >>> weights = np.array([0.1, 0.2, 0.4, 0.5], requires_grad=True)
+    >>> mt_fn(weights)
+    tensor([[ 0.25  ,  0.    , -0.0497, -0.0497],
+            [ 0.    ,  0.2475,  0.0243,  0.0243],
+            [-0.0497,  0.0243,  0.0123,  0.0123],
+            [-0.0497,  0.0243,  0.0123,  0.0123]], requires_grad=True)
+
+    This approach has the benefit of being significantly faster than the hardware-ready
+    ``metric_tensor`` function:
+
+    >>> import time
+    >>> start_time = time.process_time()
+    >>> mt = mt_fn(weights)
+    >>> time.process_time() - start_time
+    0.019
+    >>> mt_fn_2 = qml.metric_tensor(circuit)
+    >>> start_time = time.process_time()
+    >>> mt = mt_fn_2(weights)
+    >>> time.process_time() - start_time
+    0.025
+
+    This speedup becomes more drastic for larger circuits.
+    The drawback of the adjoint method is that it is only available on simulators and without
+    shot simulations.
+    """
+    if isinstance(circuit, qml.tape.QuantumTape):
+        return _adjoint_metric_tensor_tape(circuit, device)
+    if isinstance(circuit, (qml.QNode, qml.ExpvalCost)):
+        return _adjoint_metric_tensor_qnode(circuit, device, hybrid)
+
+    raise qml.QuantumFunctionError("The passed object is not a QuantumTape or QNode.")
+
+
+def _adjoint_metric_tensor_tape(tape, device):
+    """Computes the metric tensor of a tape using the adjoint method and a given device."""
+    # pylint: disable=protected-access
+    if device.shots is not None:
+        raise ValueError(
+            "The adjoint method for the metric tensor is only implemented for shots=None"
+        )
+    tape = qml.transforms.expand_trainable_multipar(tape)
+
+    # Divide all operations of a tape into trainable operations and blocks
+    # of untrainable operations after each trainable one.
+    trainable_operations, group_after_trainable_op = _group_operations(tape)
+
+    dim = 2 ** device.num_wires
+    # generate and extract initial state
+    psi = device._create_basis_state(0)
+
+    # initialize metric tensor components (which all will be real-valued)
+    like_real = qml.math.real(psi[0])
+    L = qml.math.convert_like(qml.math.zeros((tape.num_params, tape.num_params)), like_real)
+    T = qml.math.convert_like(qml.math.zeros((tape.num_params,)), like_real)
+
+    psi = _apply_operations(psi, group_after_trainable_op[-1], device)
+
+    for j, outer_op in enumerate(trainable_operations):
+        generator_1, prefactor_1 = _get_generator(outer_op)
+
+        # the state vector phi is missing a factor of 1j * prefactor_1
+        phi = device._apply_unitary(
+            psi, qml.math.convert_like(generator_1, like_real), outer_op.wires
+        )
+
+        phi_real = qml.math.reshape(qml.math.real(phi), (dim,))
+        phi_imag = qml.math.reshape(qml.math.imag(phi), (dim,))
+        diag_value = prefactor_1 ** 2 * (
+            qml.math.dot(phi_real, phi_real) + qml.math.dot(phi_imag, phi_imag)
+        )
+        L = qml.math.scatter_element_add(L, (j, j), diag_value)
+
+        lam = psi * 1.0
+        lam_real = qml.math.reshape(qml.math.real(lam), (dim,))
+        lam_imag = qml.math.reshape(qml.math.imag(lam), (dim,))
+
+        # this entry is missing a factor of 1j
+        value = prefactor_1 * (qml.math.dot(lam_real, phi_real) + qml.math.dot(lam_imag, phi_imag))
+        T = qml.math.scatter_element_add(T, (j,), value)
+
+        for i in range(j - 1, -1, -1):
+            # after first iteration of inner loop: apply U_{i+1}^\dagger
+            if i < j - 1:
+                phi = _apply_operations(phi, trainable_operations[i + 1], device, invert=True)
+            # apply V_{i}^\dagger
+            phi = _apply_operations(phi, group_after_trainable_op[i], device, invert=True)
+            lam = _apply_operations(lam, group_after_trainable_op[i], device, invert=True)
+            inner_op = trainable_operations[i]
+            # extract and apply G_i
+            generator_2, prefactor_2 = _get_generator(inner_op)
+            # this state vector is missing a factor of 1j * prefactor_2
+            mu = device._apply_unitary(lam, qml.math.convert_like(generator_2, lam), inner_op.wires)
+            phi_real = qml.math.reshape(qml.math.real(phi), (dim,))
+            phi_imag = qml.math.reshape(qml.math.imag(phi), (dim,))
+            mu_real = qml.math.reshape(qml.math.real(mu), (dim,))
+            mu_imag = qml.math.reshape(qml.math.imag(mu), (dim,))
+            # this entry is missing a factor of 1j * (-1j) = 1, i.e. none
+            value = (
+                prefactor_1
+                * prefactor_2
+                * (qml.math.dot(mu_real, phi_real) + qml.math.dot(mu_imag, phi_imag))
+            )
+            L = qml.math.scatter_element_add(
+                L, [(i, j), (j, i)], value * qml.math.convert_like(qml.math.ones((2,)), value)
+            )
+            # apply U_i^\dagger
+            lam = _apply_operations(lam, inner_op, device, invert=True)
+
+        # apply U_j and V_j
+        psi = _apply_operations(psi, [outer_op, *group_after_trainable_op[j]], device)
+
+    # postprocessing: combine L and T into the metric tensor.
+    # We require outer(conj(T), T) here, but as we skipped the factor 1j above,
+    # the stored T is real-valued. Thus we have -1j*1j*outer(T, T) = outer(T, T)
+    metric_tensor = L - qml.math.tensordot(T, T, 0)
+
+    return metric_tensor
+
+
+def _adjoint_metric_tensor_qnode(qnode, device, hybrid):
+    """Computes the metric tensor of a qnode using the adjoint method and its device.
+    For ``hybrid==True`` this wrapper accounts for classical preprocessing within the
+    QNode.
+    """
+    if device is None:
+        if isinstance(qnode, qml.ExpvalCost):
+            if qnode._multiple_devices:  # pylint: disable=protected-access
+                warnings.warn(
+                    "ExpvalCost was instantiated with multiple devices. Only the first device "
+                    "will be used to evaluate the metric tensor with the adjoint method.",
+                    UserWarning,
+                )
+            qnode = qnode.qnodes.qnodes[0]
+        device = qnode.device
+
+    cjac_fn = qml.transforms.classical_jacobian(
+        qnode, expand_fn=qml.transforms.expand_trainable_multipar
+    )
+
+    def wrapper(*args, **kwargs):
+        qnode.construct(args, kwargs)
+        mt = _adjoint_metric_tensor_tape(qnode.qtape, device)
+
+        if not hybrid:
+            return mt
+
+        cjac = cjac_fn(*args, **kwargs)
+
+        return _contract_metric_tensor_with_cjac(mt, cjac, args, qnode.interface)
+
+    return wrapper
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index 0bb985639..84092ea8e 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -259,6 +259,77 @@ def metric_tensor(tape, approx=None, allow_nonunitary=True, aux_wire=None, devic
     )
 
 
+def _contract_metric_tensor_with_cjac(mt, cjac, args, interface):
+    """Execute the contraction of pre-computed classical Jacobian(s)
+    and the metric tensor of a tape in order to obtain the hybrid
+    metric tensor of a QNode.
+
+    Args:
+        mt (array): Metric tensor of a tape (2-dimensional)
+        cjac (array or tuple[array]): The classical Jacobian of a QNode
+        args (tuple): QNode arguments
+        interface (str): QNode interface
+
+    Returns:
+        array or tuple[array]: Hybrid metric tensor(s) of the QNode.
+        The number of metric tensors depends on the number of QNode arguments
+        for which the classical Jacobian was computed, their shape on the
+        shape of these QNode arguments.
+    """
+    if isinstance(cjac, tuple):
+        # Classical processing of multiple arguments is present. Return cjac.T @ mt @ cjac
+        # as a tuple of contractions.
+        metric_tensors = tuple(
+            qml.math.tensordot(c, qml.math.tensordot(mt, c, axes=[[-1], [0]]), axes=[[0], [0]])
+            for c in cjac
+            if c is not None
+        )
+        if len(metric_tensors) == 1:
+            return metric_tensors[0]
+
+        return metric_tensors
+
+    is_square = cjac.shape == (1,) or (cjac.ndim == 2 and cjac.shape[0] == cjac.shape[1])
+
+    if is_square and qml.math.allclose(cjac, qml.numpy.eye(cjac.shape[0])):
+        # Classical Jacobian is the identity. No classical processing
+        # is present inside the QNode.
+        return mt
+
+    # TODO: Remove the following behaviour once the stacking behaviour in `qml.jacobian`
+    # has been removed. The additional arguments `args` and `interface` can be removed
+    # accordingly.
+
+    # Get number of gate arguments that were considered trainable
+    num_gate_args = qml.math.shape(mt)[0]
+    # Get trainable args
+    # pylint: disable=protected-access
+    trainable_args_idx = qml.gradients.gradient_transform._jacobian_trainable_args(args, interface)
+    # Since all arguments have the same shape, obtain shape from the first trainable arg
+    qnode_arg_shape = qml.math.shape(args[trainable_args_idx[0]])
+    num_qnode_args = len(trainable_args_idx)
+
+    if qml.math.shape(cjac) == (num_gate_args, *qnode_arg_shape):
+        # single QNode argument
+        cjac_axis = [0]
+    elif qml.math.shape(cjac) == (*qnode_arg_shape[::-1], num_gate_args, num_qnode_args):
+        # multiple QNode arguments with stacking
+        cjac_axis = [-2]
+    else:  # pragma: no cover
+        warnings.warn(
+            "Unexpected classical Jacobian encoutered, could not compute the hybrid "
+            "metric tensor of the QNode. You can still attempt to obtain the quantum "
+            "metric tensor with the `hybrid=False` parameter.",
+            UserWarning,
+        )
+        return ()
+
+    mt = qml.math.tensordot(mt, cjac, [[-1], cjac_axis])
+    mt = qml.math.tensordot(cjac, mt, [cjac_axis, [0]])
+
+    return mt
+
+
 @metric_tensor.custom_qnode_wrapper
 def qnode_execution_wrapper(self, qnode, targs, tkwargs):
     """Here, we overwrite the QNode execution wrapper in order
@@ -311,32 +382,7 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
         kwargs.pop("shots", False)
         cjac = cjac_fn(*args, **kwargs)
 
-        if isinstance(cjac, tuple):
-            if len(cjac) == 1:
-                cjac = cjac[0]
-            else:
-                # Classical processing of multiple arguments is present. Return cjac.T @ mt @ cjac.
-                metric_tensors = []
-
-                for c in cjac:
-                    if c is not None:
-                        _mt = qml.math.tensordot(mt, c, axes=[[-1], [0]])
-                        _mt = qml.math.tensordot(c, _mt, axes=[[0], [0]])
-                        metric_tensors.append(_mt)
-
-                return tuple(metric_tensors)
-
-        is_square = cjac.shape == (1,) or (cjac.ndim == 2 and cjac.shape[0] == cjac.shape[1])
-
-        if is_square and qml.math.allclose(cjac, qml.numpy.eye(cjac.shape[0])):
-            # Classical Jacobian is the identity. No classical processing
-            # is present inside the QNode.
-            return mt
-
-        # Classical processing of a single argument is present. Return mt @ cjac.
-        mt = qml.math.tensordot(mt, cjac, [[-1], [0]])
-        mt = qml.math.tensordot(cjac, mt, [[0], [0]])
-        return mt
+        return _contract_metric_tensor_with_cjac(mt, cjac, args, qnode.interface)
 
     return wrapper
 
@@ -561,7 +607,6 @@ def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
     # and non-trainable parameter indices
     graph = tape.graph
     par_idx_to_trainable_idx = {idx: i for i, idx in enumerate(sorted(tape.trainable_params))}
-    layers = graph.iterate_parametrized_layers()
     layers = [
         LayerData(
             layer.pre_ops,
@@ -569,8 +614,10 @@ def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
             tuple(par_idx_to_trainable_idx[idx] for idx in layer[2]),
             layer.post_ops,
         )
-        for layer in layers
+        for layer in graph.iterate_parametrized_layers()
     ]
+    if len(layers) <= 1:
+        return diag_tapes, diag_proc_fn
 
     # Get default for aux_wire
     aux_wire = _get_aux_wire(aux_wire, tape, device_wires)
diff --git a/pennylane/transforms/tape_expand.py b/pennylane/transforms/tape_expand.py
index f0c42de35..6ef54d331 100644
--- a/pennylane/transforms/tape_expand.py
+++ b/pennylane/transforms/tape_expand.py
@@ -96,14 +96,14 @@ def create_expand_fn(depth, stop_at=None, device=None, docstring=None):
         else:
             stop_at &= device.stopping_condition
 
-    def expand_fn(tape, _depth=depth, **kwargs):
+    def expand_fn(tape, depth=depth, **kwargs):
 
         with qml.tape.stop_recording():
 
             if stop_at is None:
-                tape = tape.expand(depth=_depth)
+                tape = tape.expand(depth=depth)
             elif not all(stop_at(op) for op in tape.operations):
-                tape = tape.expand(depth=_depth, stop_at=stop_at)
+                tape = tape.expand(depth=depth, stop_at=stop_at)
             else:
                 return tape
 
@@ -140,6 +140,29 @@ expand_multipar = create_expand_fn(
     docstring=_expand_multipar_doc,
 )
 
+_expand_trainable_multipar_doc = """Expand out a tape so that all its trainable
+operations have a single parameter.
+
+This is achieved by decomposing all trainable operations that do not have
+a generator, up to maximum depth ``depth``.
+For a sufficient ``depth``, it should always be possible to obtain a tape containing
+only single-parameter operations.
+
+Args:
+    tape (.QuantumTape): the input tape to expand
+    depth (int) : the maximum expansion depth
+    **kwargs: additional keyword arguments are ignored
+
+Returns:
+    .QuantumTape: the expanded tape
+"""
+
+expand_trainable_multipar = create_expand_fn(
+    depth=10,
+    stop_at=not_tape | is_measurement | has_nopar | (~is_trainable) | has_gen,
+    docstring=_expand_trainable_multipar_doc,
+)
+
 
 _expand_nonunitary_gen_doc = """Expand out a tape so that all its parametrized
 operations have a unitary generator.


Test patch:
diff --git a/tests/test_vqe.py b/tests/test_vqe.py
index 2d731f7af..ce05eb337 100644
--- a/tests/test_vqe.py
+++ b/tests/test_vqe.py
@@ -591,7 +591,7 @@ class TestVQE:
         """Test that the metric tensor can be calculated."""
 
         dev = qml.device("default.qubit", wires=3)
-        p = np.array([1.0, 1.0, 1.0])
+        p = pnp.array([1.0, 1.0, 1.0], requires_grad=True)
 
         def ansatz(params, **kwargs):
             qml.RX(params[0], wires=0)
@@ -603,7 +603,7 @@ class TestVQE:
         qnodes = qml.ExpvalCost(ansatz, h, dev)
         mt = qml.metric_tensor(qnodes, approx=approx)(p)
         assert mt.shape == (3, 3)
-        assert isinstance(mt, np.ndarray)
+        assert isinstance(mt, pnp.ndarray)
 
     def test_multiple_devices(self, mocker):
         """Test that passing multiple devices to ExpvalCost works correctly"""
@@ -618,6 +618,7 @@ class TestVQE:
         qnodes = qml.ExpvalCost(qml.templates.BasicEntanglerLayers, h, dev)
         np.random.seed(1967)
         w = np.random.random(qml.templates.BasicEntanglerLayers.shape(n_layers=3, n_wires=2))
+        w = pnp.array(w, requires_grad=True)
 
         res = qnodes(w)
 
diff --git a/tests/transforms/test_adjoint_metric_tensor.py b/tests/transforms/test_adjoint_metric_tensor.py
new file mode 100644
index 000000000..3903dfa1e
--- /dev/null
+++ b/tests/transforms/test_adjoint_metric_tensor.py
@@ -0,0 +1,861 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for the adjoint_metric_tensor function.
+"""
+import pytest
+from pennylane import numpy as np
+import pennylane as qml
+
+from pennylane.transforms.adjoint_metric_tensor import _apply_operations
+
+
+class TestApplyOperations:
+    """Tests the application of operations via the helper function
+    _apply_operations used in the adjoint metric tensor."""
+
+    device = qml.device("default.qubit", wires=2)
+    x = 0.5
+
+    def test_simple_operation(self):
+        """Test that an operation is applied correctly."""
+        op = qml.RX(self.x, wires=0)
+        out = _apply_operations(self.device._state, op, self.device)
+        out = qml.math.reshape(out, 4)
+        exp = np.array([np.cos(self.x / 2), 0.0, -1j * np.sin(self.x / 2), 0.0])
+        assert np.allclose(out, exp)
+        assert not op.inverse
+
+    def test_simple_operation_inv(self):
+        """Test that an operation is applied correctly when using invert=True
+        but does not alter the operation (in particular its inverse flag) that is used."""
+        op = qml.RX(self.x, wires=0)
+        out = _apply_operations(self.device._state, op, self.device, invert=True)
+        out = qml.math.reshape(out, 4)
+        exp = np.array([np.cos(self.x / 2), 0.0, 1j * np.sin(self.x / 2), 0.0])
+        assert np.allclose(out, exp)
+        assert not op.inverse
+
+    def test_inv_operation(self):
+        """Test that an operation with inverse=True is applied correctly
+        but does not alter the operation (in particular its inverse flag) that is used."""
+        op = qml.RX(self.x, wires=0).inv()
+        out = _apply_operations(self.device._state, op, self.device)
+        out = qml.math.reshape(out, 4)
+        exp = np.array([np.cos(self.x / 2), 0.0, 1j * np.sin(self.x / 2), 0.0])
+        assert np.allclose(out, exp)
+        assert op.inverse
+
+    def test_inv_operation_inv(self):
+        """Test that an operation with inverse=True is applied correctly when using invert=True
+        but does not alter the operation (in particular its inverse flag) that is used."""
+        op = qml.RX(self.x, wires=0).inv()
+        out = _apply_operations(self.device._state, op, self.device, invert=True)
+        out = qml.math.reshape(out, 4)
+        exp = np.array([np.cos(self.x / 2), 0.0, -1j * np.sin(self.x / 2), 0.0])
+        assert np.allclose(out, exp)
+        assert op.inverse
+
+    def test_operation_group(self):
+        """Test that a group of operations with is applied correctly
+        but does not alter the operations (in particular their order and
+        inverse flags) that are used."""
+        op = [qml.RX(self.x, wires=0).inv(), qml.Hadamard(wires=1), qml.CNOT(wires=[1, 0])]
+        out = _apply_operations(self.device._state, op, self.device)
+        out = qml.math.reshape(out, 4)
+        exp = np.array(
+            [
+                np.cos(self.x / 2) / np.sqrt(2),
+                1j * np.sin(self.x / 2) / np.sqrt(2),
+                1j * np.sin(self.x / 2) / np.sqrt(2),
+                np.cos(self.x / 2) / np.sqrt(2),
+            ]
+        )
+        assert np.allclose(out, exp)
+        assert isinstance(op[0], qml.RX) and op[0].inverse
+        assert isinstance(op[1], qml.Hadamard) and not op[1].inverse
+        assert isinstance(op[2], qml.CNOT) and not op[2].inverse
+
+    def test_operation_group_inv(self):
+        """Test that a group of operations with is applied correctly when using invert=True
+        but does not alter the operations (in particular their order and
+        inverse flags) that are used."""
+        op = [qml.RX(self.x, wires=0).inv(), qml.Hadamard(wires=1), qml.CNOT(wires=[1, 0])]
+        out = _apply_operations(self.device._state, op, self.device, invert=True)
+        out = qml.math.reshape(out, 4)
+        exp = np.array(
+            [
+                np.cos(self.x / 2) / np.sqrt(2),
+                np.cos(self.x / 2) / np.sqrt(2),
+                -1j * np.sin(self.x / 2) / np.sqrt(2),
+                -1j * np.sin(self.x / 2) / np.sqrt(2),
+            ]
+        )
+        assert np.allclose(out, exp)
+        assert isinstance(op[0], qml.RX) and op[0].inverse
+        assert isinstance(op[1], qml.Hadamard) and not op[1].inverse
+        assert isinstance(op[2], qml.CNOT) and not op[2].inverse
+
+    def test_qubit_statevector(self):
+        """Test that a statevector preparation is applied correctly."""
+        state = np.array([0.4, 1.2 - 0.2j, 9.5, -0.3 + 1.1j])
+        state /= np.linalg.norm(state, ord=2)
+        op = qml.QubitStateVector(state, wires=self.device.wires)
+        out = _apply_operations(None, op, self.device, invert=False)
+        out = qml.math.reshape(out, 4)
+        assert np.allclose(out, state)
+
+    def test_error_qubit_statevector(self):
+        """Test that an error is raised for a statevector preparation with invert=True."""
+        state = np.array([0.4, 1.2 - 0.2j, 9.5, -0.3 + 1.1j])
+        state = np.array([0.4, 1.2 - 0.2j, 9.5, -0.3 + 1.1j])
+        state /= np.linalg.norm(state, ord=2)
+        op = qml.QubitStateVector(state, wires=self.device.wires)
+        with pytest.raises(ValueError, match="Can't invert state preparation."):
+            _apply_operations(None, op, self.device, invert=True)
+
+    def test_basisstate(self):
+        """Test that a basis state preparation is applied correctly."""
+        op = qml.BasisState(np.array([1, 0]), wires=self.device.wires)
+        out = _apply_operations(None, op, self.device, invert=False)
+        out = qml.math.reshape(out, 4)
+        exp = np.array([0.0, 0.0, 1.0, 0.0])
+        assert np.allclose(out, exp)
+
+    def test_error_basisstate(self):
+        """Test that an error is raised for a basis state preparation with invert=True."""
+        op = qml.BasisState(np.array([1, 0]), wires=self.device.wires)
+        with pytest.raises(ValueError, match="Can't invert state preparation."):
+            _apply_operations(None, op, self.device, invert=True)
+
+
+@pytest.mark.parametrize("invert", [False, True])
+class TestApplyOperationsDifferentiability:
+    """Tests the differentiability of applying operations via the helper function
+    _apply_operations used in the adjoint metric tensor."""
+
+    x = 0.5
+
+    def test_simple_operation_autograd(self, invert):
+        """Test differentiability for a simple operation with Autograd."""
+        device = qml.device("default.qubit.autograd", wires=2)
+        x = np.array(self.x, requires_grad=True)
+        r_fn = lambda x: qml.math.real(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        i_fn = lambda x: qml.math.imag(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        out = qml.jacobian(r_fn)(x) + 1j * qml.jacobian(i_fn)(x)
+        exp = (
+            np.array([[-np.sin(self.x / 2), 0.0], [-1j * (-1) ** invert * np.cos(self.x / 2), 0.0]])
+            / 2
+        )
+        assert np.allclose(out, exp)
+
+    def test_simple_operation_jax(self, invert):
+        """Test differentiability for a simple operation with JAX."""
+        jax = pytest.importorskip("jax")
+        device = qml.device("default.qubit.jax", wires=2)
+        x = jax.numpy.array(self.x)
+        r_fn = lambda x: qml.math.real(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        i_fn = lambda x: qml.math.imag(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        out = jax.jacobian(r_fn)(x) + 1j * jax.jacobian(i_fn)(x)
+        exp = (
+            np.array([[-np.sin(self.x / 2), 0.0], [-1j * (-1) ** invert * np.cos(self.x / 2), 0.0]])
+            / 2
+        )
+        assert np.allclose(out, exp)
+
+    def test_simple_operation_tf(self, invert):
+        """Test differentiability for a simple operation with TensorFlow."""
+        tf = pytest.importorskip("tensorflow")
+        device = qml.device("default.qubit.tf", wires=2)
+        x = tf.Variable(self.x, dtype=tf.float64)
+        r_fn = lambda x: qml.math.real(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        i_fn = lambda x: qml.math.imag(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        with tf.GradientTape(persistent=True) as tape:
+            r_state = r_fn(x)
+            i_state = i_fn(x)
+        out = qml.math.complex(tape.jacobian(r_state, x), tape.jacobian(i_state, x))
+        exp = (
+            np.array([[-np.sin(self.x / 2), 0.0], [-1j * (-1) ** invert * np.cos(self.x / 2), 0.0]])
+            / 2
+        )
+        assert np.allclose(out, exp)
+
+    def test_simple_operation_torch(self, invert):
+        """Test differentiability for a simple operation with Torch."""
+        torch = pytest.importorskip("torch")
+        jac_fn = torch.autograd.functional.jacobian
+        device = qml.device("default.qubit.torch", wires=2)
+        x = torch.tensor(self.x, requires_grad=True)
+        r_fn = lambda x: qml.math.real(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        i_fn = lambda x: qml.math.imag(
+            _apply_operations(device._state, qml.RX(x, wires=0), device, invert)
+        )
+        out = jac_fn(r_fn, x) + 1j * jac_fn(i_fn, x)
+        exp = (
+            np.array([[-np.sin(self.x / 2), 0.0], [-1j * (-1) ** invert * np.cos(self.x / 2), 0.0]])
+            / 2
+        )
+        assert np.allclose(out, exp)
+
+
+fixed_pars = [-0.2, 0.2, 0.5, 0.3, 0.7]
+
+
+def fubini_ansatz0(params, wires=None):
+    qml.RX(params[0], wires=0)
+    qml.RY(fixed_pars[0], wires=0)
+    qml.CNOT(wires=[wires[0], wires[1]])
+    qml.RZ(params[1], wires=0)
+    qml.CNOT(wires=[wires[0], wires[1]])
+
+
+def fubini_ansatz1(params, wires=None):
+    qml.RX(fixed_pars[1], wires=0)
+    for wire in wires:
+        qml.Rot(*params[0][wire], wires=wire)
+    qml.CNOT(wires=[0, 1])
+    qml.RY(fixed_pars[1], wires=0).inv()
+    qml.CNOT(wires=[1, 2])
+    for wire in wires:
+        qml.Rot(*params[1][wire], wires=wire)
+    qml.CNOT(wires=[1, 2])
+    qml.RX(fixed_pars[2], wires=1)
+
+
+def fubini_ansatz2(params, wires=None):
+    params0 = params[0]
+    params1 = params[1]
+    qml.RX(fixed_pars[1], wires=0)
+    qml.Rot(*fixed_pars[2:5], wires=1)
+    qml.CNOT(wires=[0, 1])
+    qml.RY(params0, wires=0)
+    qml.RY(params0, wires=1)
+    qml.CNOT(wires=[0, 1])
+    qml.RX(params1, wires=0).inv()
+    qml.RX(params1, wires=1)
+
+
+def fubini_ansatz3(params, wires=None):
+    params0 = params[0]
+    params1 = params[1]
+    params2 = params[2]
+    qml.RX(fixed_pars[1], wires=0)
+    qml.RX(fixed_pars[3], wires=1)
+    qml.CNOT(wires=[0, 1])
+    qml.CNOT(wires=[1, 2])
+    qml.RX(params0, wires=0)
+    qml.RX(params0, wires=1)
+    qml.CNOT(wires=[0, 1])
+    qml.CNOT(wires=[1, 2])
+    qml.CNOT(wires=[2, 0])
+    qml.RY(params1, wires=0)
+    qml.RY(params1, wires=1)
+    qml.RY(params1, wires=2)
+    qml.RZ(params2, wires=0)
+    qml.RZ(params2, wires=1)
+    qml.RZ(params2, wires=2)
+
+
+def fubini_ansatz4(params00, params_rest, wires=None):
+    params01 = params_rest[0]
+    params10 = params_rest[1]
+    params11 = params_rest[2]
+    qml.RY(fixed_pars[3], wires=0)
+    qml.RY(fixed_pars[2], wires=1)
+    qml.CNOT(wires=[0, 1])
+    qml.CNOT(wires=[1, 2])
+    qml.RY(fixed_pars[4], wires=0)
+    qml.RX(params00, wires=0)
+    qml.CNOT(wires=[0, 1])
+    qml.RX(params01, wires=1)
+    qml.RZ(params10, wires=1)
+    qml.CNOT(wires=[0, 1])
+    qml.RZ(params11, wires=1)
+
+
+def fubini_ansatz5(params, wires=None):
+    fubini_ansatz4(params[0], [params[0], params[1], params[1]], wires=wires)
+
+
+def fubini_ansatz6(params, wires=None):
+    fubini_ansatz4(params[0], [params[0], params[1], -params[1]], wires=wires)
+
+
+def fubini_ansatz7(params0, params1, wires=None):
+    fubini_ansatz4(params0[0], [params0[1], params1[0], params1[1]], wires=wires)
+
+
+def fubini_ansatz8(x, wires=None):
+    qml.RX(fixed_pars[0], wires=0)
+    qml.RX(x, wires=0)
+
+
+def fubini_ansatz9(params, wires=None):
+    params0 = params[0]
+    params1 = params[1]
+    qml.RX(fixed_pars[1], wires=[0])
+    qml.RY(fixed_pars[3], wires=[0])
+    qml.RZ(fixed_pars[2], wires=[0])
+    qml.RX(fixed_pars[2], wires=[1])
+    qml.RY(fixed_pars[2], wires=[1])
+    qml.RZ(fixed_pars[4], wires=[1])
+    qml.CNOT(wires=[0, 1])
+    qml.RX(fixed_pars[0], wires=[0])
+    qml.RY(fixed_pars[1], wires=[0])
+    qml.RZ(fixed_pars[3], wires=[0])
+    qml.RX(fixed_pars[1], wires=[1])
+    qml.RY(fixed_pars[2], wires=[1])
+    qml.RZ(fixed_pars[0], wires=[1])
+    qml.CNOT(wires=[0, 1])
+    qml.RX(params0, wires=[0])
+    qml.RX(params0, wires=[1])
+    qml.CNOT(wires=[0, 1])
+    qml.RY(fixed_pars[4], wires=[1])
+    qml.RY(params1, wires=[0])
+    qml.RY(params1, wires=[1])
+    qml.CNOT(wires=[0, 1])
+    qml.RX(fixed_pars[2], wires=[1])
+
+
+def fubini_ansatz10(weights, wires=None):
+    qml.templates.BasicEntanglerLayers(weights, wires=[0, 1])
+
+
+B = np.array(
+    [
+        [
+            [0.73, 0.49, 0.04],
+            [0.29, 0.45, 0.59],
+            [0.64, 0.06, 0.26],
+        ],
+        [
+            [0.93, 0.14, 0.46],
+            [0.31, 0.83, 0.79],
+            [0.25, 0.40, 0.16],
+        ],
+    ],
+    requires_grad=True,
+)
+fubini_ansatze_tape = [fubini_ansatz0, fubini_ansatz1, fubini_ansatz8]
+fubini_params_tape = [
+    (np.array([0.3434, -0.7245345], requires_grad=True),),
+    (B,),
+    (np.array(-0.1735, requires_grad=True),),
+]
+
+fubini_ansatze = [
+    fubini_ansatz0,
+    fubini_ansatz1,
+    fubini_ansatz2,
+    fubini_ansatz3,
+    fubini_ansatz4,
+    fubini_ansatz5,
+    fubini_ansatz6,
+    fubini_ansatz7,
+    fubini_ansatz8,
+    fubini_ansatz9,
+    fubini_ansatz10,
+]
+
+fubini_params = [
+    (np.array([0.3434, -0.7245345], requires_grad=True),),
+    (B,),
+    (np.array([-0.1111, -0.2222], requires_grad=True),),
+    (np.array([-0.1111, -0.2222, 0.4554], requires_grad=True),),
+    (
+        np.array(-0.1735, requires_grad=True),
+        np.array([-0.1735, -0.2846, -0.2846], requires_grad=True),
+    ),
+    (np.array([-0.1735, -0.2846], requires_grad=True),),
+    (np.array([-0.1735, -0.2846], requires_grad=True),),
+    (
+        np.array([-0.1735, -0.2846], requires_grad=True),
+        np.array([0.9812, -0.1492], requires_grad=True),
+    ),
+    (np.array(-0.1735, requires_grad=True),),
+    (np.array([-0.1111, 0.3333], requires_grad=True),),
+    (np.array([[0.21, 9.29], [-0.2, 0.12], [0.3, -2.1]], requires_grad=True),),
+]
+
+
+def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
+    """Compute the metric tensor by full state vector
+    differentiation via autograd."""
+    dev = qml.device("default.qubit", wires=num_wires)
+
+    @qml.qnode(dev)
+    def qnode(*params):
+        ansatz(*params, wires=dev.wires)
+        return qml.state()
+
+    def mt(*params):
+        state = qnode(*params)
+        rqnode = lambda *params: np.real(qnode(*params))
+        iqnode = lambda *params: np.imag(qnode(*params))
+        rjac = qml.jacobian(rqnode)(*params)
+        ijac = qml.jacobian(iqnode)(*params)
+
+        if not mimic_autograd:
+            if len(params) > 1 and all(
+                qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:]
+            ):
+                rjac = tuple(qml.math.transpose(rjac))
+                ijac = tuple(qml.math.transpose(ijac))
+
+        if isinstance(rjac, tuple):
+            out = []
+            for rc, ic in zip(rjac, ijac):
+                c = rc + 1j * ic
+                psidpsi = np.tensordot(np.conj(state), c, axes=([0], [0]))
+                out.append(
+                    np.real(
+                        np.tensordot(np.conj(c), c, axes=([0], [0]))
+                        - np.tensordot(np.conj(psidpsi), psidpsi, axes=0)
+                    )
+                )
+            return tuple(out)
+
+        if (
+            mimic_autograd
+            and len(params) > 1
+            and all(qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:])
+        ):
+            jac_contract = [-2]
+        else:
+            jac_contract = [0]
+        jac = rjac + 1j * ijac
+        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], jac_contract))
+        return np.real(
+            np.tensordot(np.conj(jac), jac, axes=(jac_contract, jac_contract))
+            - np.tensordot(np.conj(psidpsi), psidpsi, axes=0)
+        )
+
+    return mt
+
+
+@pytest.mark.parametrize("ansatz, params", list(zip(fubini_ansatze_tape, fubini_params_tape)))
+class TestAdjointMetricTensorTape:
+    """Test the adjoint method for the metric tensor when calling it directly on
+    a tape.
+    """
+
+    num_wires = 3
+
+    def test_correct_output_tape_autograd(self, ansatz, params):
+        """Test that the output is correct when using Autograd and
+        calling the adjoint metric tensor directly on a tape."""
+        expected = autodiff_metric_tensor(ansatz, self.num_wires, True)(*params)
+        dev = qml.device("default.qubit.autograd", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        circuit(*params)
+        mt = qml.adjoint_metric_tensor(circuit.qtape, dev)
+        expected = qml.math.reshape(expected, qml.math.shape(mt))
+        assert qml.math.allclose(mt, expected)
+
+        mt = qml.adjoint_metric_tensor(circuit, hybrid=False)(*params)
+        assert qml.math.allclose(mt, expected)
+
+    @pytest.mark.skip("JAX does not support forward pass executiong of the metric tensor.")
+    def test_correct_output_tape_jax(self, ansatz, params):
+        """Test that the output is correct when using JAX and
+        calling the adjoint metric tensor directly on a tape."""
+
+        jax = pytest.importorskip("jax")
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        j_params = tuple(jax.numpy.array(p) for p in params)
+        dev = qml.device("default.qubit.jax", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        circuit(*j_params)
+        mt = qml.adjoint_metric_tensor(circuit.qtape, dev)
+        expected = qml.math.reshape(expected, qml.math.shape(mt))
+        assert qml.math.allclose(mt, expected)
+
+        mt = qml.adjoint_metric_tensor(circuit, hybrid=False)(*j_params)
+        assert qml.math.allclose(mt, expected)
+
+    def test_correct_output_tape_torch(self, ansatz, params):
+        """Test that the output is correct when using Torch and
+        calling the adjoint metric tensor directly on a tape."""
+
+        torch = pytest.importorskip("torch")
+
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        t_params = tuple(torch.tensor(p, requires_grad=True) for p in params)
+        dev = qml.device("default.qubit.torch", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        circuit(*t_params)
+        mt = qml.adjoint_metric_tensor(circuit.qtape, dev)
+        expected = qml.math.reshape(expected, qml.math.shape(mt))
+        assert qml.math.allclose(mt.detach().numpy(), expected)
+
+        mt = qml.adjoint_metric_tensor(circuit, hybrid=False)(*t_params)
+        assert qml.math.allclose(mt, expected)
+
+    def test_correct_output_tape_tf(self, ansatz, params):
+        """Test that the output is correct when using TensorFlow and
+        calling the adjoint metric tensor directly on a tape."""
+
+        tf = pytest.importorskip("tensorflow")
+
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        t_params = tuple(tf.Variable(p) for p in params)
+        dev = qml.device("default.qubit.tf", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        with tf.GradientTape() as t:
+            circuit(*t_params)
+            mt = qml.adjoint_metric_tensor(circuit.qtape, dev)
+
+        expected = qml.math.reshape(expected, qml.math.shape(mt))
+        assert qml.math.allclose(mt, expected)
+
+        with tf.GradientTape() as t:
+            mt = qml.adjoint_metric_tensor(circuit, hybrid=False)(*t_params)
+        assert qml.math.allclose(mt, expected)
+
+
+class TestAdjointMetricTensorQNode:
+    """Test the adjoint method for the metric tensor when calling it on
+    a QNode.
+    """
+
+    num_wires = 3
+
+    @pytest.mark.parametrize("ansatz, params", list(zip(fubini_ansatze, fubini_params)))
+    def test_correct_output_qnode_autograd(self, ansatz, params):
+        """Test that the output is correct when using Autograd and
+        calling the adjoint metric tensor on a QNode."""
+        expected = autodiff_metric_tensor(ansatz, self.num_wires, True)(*params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        mt = qml.adjoint_metric_tensor(circuit)(*params)
+
+        print(mt)
+        print(expected)
+        if isinstance(mt, tuple):
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt, expected))
+        else:
+            assert qml.math.allclose(mt, expected)
+
+    @pytest.mark.skip("JAX does not support forward pass executiong of the metric tensor.")
+    @pytest.mark.parametrize("ansatz, params", list(zip(fubini_ansatze, fubini_params)))
+    def test_correct_output_qnode_jax(self, ansatz, params):
+        """Test that the output is correct when using JAX and
+        calling the adjoint metric tensor on a QNode."""
+
+        jax = pytest.importorskip("jax")
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        j_params = tuple(jax.numpy.array(p) for p in params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        mt = qml.adjoint_metric_tensor(circuit)(*j_params)
+
+        if isinstance(mt, tuple):
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt, expected))
+        else:
+            assert qml.math.allclose(mt, expected)
+
+    @pytest.mark.parametrize("ansatz, params", list(zip(fubini_ansatze, fubini_params)))
+    def test_correct_output_qnode_torch(self, ansatz, params):
+        """Test that the output is correct when using Torch and
+        calling the adjoint metric tensor on a QNode."""
+
+        torch = pytest.importorskip("torch")
+
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        t_params = tuple(torch.tensor(p, requires_grad=True, dtype=torch.float64) for p in params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        mt = qml.adjoint_metric_tensor(circuit)(*t_params)
+
+        print(mt)
+        print(expected)
+        if isinstance(mt, tuple):
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt, expected))
+        else:
+            assert qml.math.allclose(mt, expected)
+
+    @pytest.mark.parametrize("ansatz, params", list(zip(fubini_ansatze, fubini_params)))
+    def test_correct_output_qnode_tf(self, ansatz, params):
+        """Test that the output is correct when using TensorFlow and
+        calling the adjoint metric tensor on a QNode."""
+
+        tf = pytest.importorskip("tensorflow")
+
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        t_params = tuple(tf.Variable(p, dtype=tf.float64) for p in params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        with tf.GradientTape() as t:
+            mt = qml.adjoint_metric_tensor(circuit)(*t_params)
+
+        if isinstance(mt, tuple):
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt, expected))
+        else:
+            assert qml.math.allclose(mt, expected)
+
+    def test_autograd_with_other_device(self):
+        """Test passing an extra device to the QNode wrapper."""
+        ansatz = fubini_ansatz2
+        params = fubini_params[2]
+
+        exp_fn = autodiff_metric_tensor(ansatz, self.num_wires, True)
+        expected = qml.jacobian(exp_fn)(*params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+        dev2 = qml.device("default.qubit.autograd", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        mt = qml.jacobian(qml.adjoint_metric_tensor(circuit, device=dev2))(*params)
+
+        if isinstance(mt, tuple):
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt, expected))
+        else:
+            assert qml.math.allclose(mt, expected)
+
+
+diff_fubini_ansatze = [
+    fubini_ansatz0,
+    fubini_ansatz2,
+    fubini_ansatz10,
+]
+
+diff_fubini_params = [
+    fubini_params[0],
+    fubini_params[2],
+    fubini_params[10],
+]
+
+
+@pytest.mark.parametrize("ansatz, params", list(zip(diff_fubini_ansatze, diff_fubini_params)))
+class TestAdjointMetricTensorDifferentiability:
+    """Test the differentiability of the adjoint method for the metric
+    tensor when calling it on a QNode.
+    """
+
+    num_wires = 3
+
+    def test_autograd(self, ansatz, params):
+        """Test that the derivative is correct when using Autograd and
+        calling the adjoint metric tensor on a QNode."""
+        exp_fn = autodiff_metric_tensor(ansatz, self.num_wires, True)
+        expected = qml.jacobian(exp_fn)(*params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        mt_jac = qml.jacobian(qml.adjoint_metric_tensor(circuit))(*params)
+
+        if isinstance(mt_jac, tuple):
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt_jac, expected))
+        else:
+            assert qml.math.allclose(mt_jac, expected)
+
+    def test_correct_output_qnode_jax(self, ansatz, params):
+        """Test that the derivative is correct when using JAX and
+        calling the adjoint metric tensor on a QNode."""
+
+        jax = pytest.importorskip("jax")
+        from jax.config import config
+
+        config.update("jax_enable_x64", True)
+
+        expected = qml.jacobian(autodiff_metric_tensor(ansatz, self.num_wires))(*params)
+        j_params = tuple(jax.numpy.array(p) for p in params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        mt_fn = qml.adjoint_metric_tensor(circuit, hybrid=True)
+        argnums = list(range(len(params)))
+        mt_jac = jax.jacobian(mt_fn, argnums=argnums)(*j_params)
+
+        if isinstance(mt_jac, tuple):
+            if not isinstance(expected, tuple) and len(mt_jac) == 1:
+                expected = (expected,)
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt_jac, expected))
+        else:
+            assert qml.math.allclose(mt_jac, expected)
+
+    def test_correct_output_qnode_torch(self, ansatz, params):
+        """Test that the derivative is correct when using Torch and
+        calling the adjoint metric tensor on a QNode."""
+
+        torch = pytest.importorskip("torch")
+
+        expected = qml.jacobian(autodiff_metric_tensor(ansatz, self.num_wires))(*params)
+        t_params = tuple(torch.tensor(p, requires_grad=True, dtype=torch.float64) for p in params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        mt_fn = qml.adjoint_metric_tensor(circuit)
+        mt_jac = torch.autograd.functional.jacobian(mt_fn, *t_params)
+
+        if isinstance(mt_jac, tuple):
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt_jac, expected))
+        else:
+            assert qml.math.allclose(mt_jac, expected)
+
+    def test_correct_output_qnode_tf(self, ansatz, params):
+        """Test that the derivative is correct when using TensorFlow and
+        calling the adjoint metric tensor on a QNode."""
+
+        tf = pytest.importorskip("tensorflow")
+
+        expected = qml.jacobian(autodiff_metric_tensor(ansatz, self.num_wires))(*params)
+        t_params = tuple(tf.Variable(p, dtype=tf.float64) for p in params)
+        dev = qml.device("default.qubit", wires=self.num_wires)
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires)
+            return qml.expval(qml.PauliZ(0))
+
+        with tf.GradientTape() as t:
+            mt = qml.adjoint_metric_tensor(circuit)(*t_params)
+
+        mt_jac = t.jacobian(mt, t_params)
+        if isinstance(mt_jac, tuple):
+            if not isinstance(expected, tuple) and len(mt_jac) == 1:
+                expected = (expected,)
+            assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt_jac, expected))
+        else:
+            assert qml.math.allclose(mt_jac, expected)
+
+
+class TestErrors:
+    """Test that errors are raised correctly."""
+
+    def test_error_wrong_object_passed(self):
+        """Test that an error is raised if neither a tape nor a QNode is passed."""
+
+        def ansatz(x, y):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        with pytest.raises(qml.QuantumFunctionError, match="The passed object is not a "):
+            qml.adjoint_metric_tensor(ansatz, device=dev)
+
+    def test_error_finite_shots(self):
+        """Test that an error is raised if the device has a finite number of shots set."""
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(0.2, wires=0)
+            qml.RY(1.9, wires=1)
+        dev = qml.device("default.qubit", wires=2, shots=1)
+
+        with pytest.raises(ValueError, match="The adjoint method for the metric tensor"):
+            qml.adjoint_metric_tensor(tape, device=dev)
+
+    def test_warning_multiple_devices(self):
+        """Test that a warning is issued if an ExpvalCost with multiple
+        devices is passed."""
+        dev1 = qml.device("default.qubit", wires=2)
+        dev2 = qml.device("default.qubit", wires=1)
+        H = qml.Hamiltonian([0.2, 0.9], [qml.PauliZ(0), qml.PauliY(0)])
+
+        def ansatz(x, wires):
+            qml.RX(x, wires=0)
+
+        cost = qml.ExpvalCost(ansatz, H, [dev1, dev2])
+        with pytest.warns(UserWarning, match="ExpvalCost was instantiated"):
+            mt = qml.adjoint_metric_tensor(cost)
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index 999a20db3..a97d8556f 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -63,15 +63,15 @@ class TestMetricTensor:
         """Test that the MultiRZ gate is correctly decomposed"""
         dev = qml.device("default.qubit", wires=3)
 
+        @qml.qnode(dev, diff_method=diff_method)
         def circuit(a, b):
             qml.RX(a, wires=0)
             qml.MultiRZ(b, wires=[0, 1, 2])
             return qml.expval(qml.PauliX(0))
 
-        circuit = qml.QNode(circuit, dev, diff_method=diff_method)
-        params = [0.1, 0.2]
+        params = np.array([0.1, 0.2], requires_grad=True)
         result = qml.metric_tensor(circuit, approx="block-diag")(*params)
-        assert result.shape == (2, 2)
+        assert qml.math.shape(result) == (2, 2)
 
     @pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop"])
     def test_parameter_fan_out(self, diff_method):
@@ -216,9 +216,9 @@ class TestMetricTensor:
 
         circuit = qml.QNode(circuit, dev)
 
-        a = 0.432
-        b = 0.12
-        c = -0.432
+        a = np.array(0.432, requires_grad=True)
+        b = np.array(0.12, requires_grad=True)
+        c = np.array(-0.432, requires_grad=True)
 
         # evaluate metric tensor
         g_diag = qml.metric_tensor(circuit, approx="diag")(a, b, c)
@@ -306,7 +306,8 @@ class TestMetricTensor:
         b = 0.1
         c = 0.5
 
-        def final(x, y, z, h, g, f):
+        def final(params):
+            x, y, z, h, g, f = params
             non_parametrized_layer(a, b, c)
             qml.RX(x, wires=0)
             qml.RY(-y, wires=1).inv()
@@ -327,10 +328,12 @@ class TestMetricTensor:
         computation."""
         dev, circuit, non_parametrized_layer, a, b, c = sample_circuit
 
-        params = [-0.282203, 0.145554, 0.331624, -0.163907, 0.57662, 0.081272]
-        x, y, z, h, g, f = params
+        params = np.array(
+            [-0.282203, 0.145554, 0.331624, -0.163907, 0.57662, 0.081272],
+            requires_grad=True,
+        )
 
-        G = qml.metric_tensor(circuit, approx="block-diag")(*params)
+        G = qml.metric_tensor(circuit, approx="block-diag")(params)
 
         # ============================================
         # Test block-diag metric tensor of first layer is correct.
@@ -388,7 +391,8 @@ class TestMetricTensor:
         #   qml.RZ(g, wires=2)
         G2 = np.zeros([2, 2])
 
-        def layer2_diag(x, y, z, h, g, f):
+        def layer2_diag(params):
+            x, y, z, h, g, f = params
             non_parametrized_layer(a, b, c)
             qml.RX(x, wires=0)
             qml.RY(y, wires=1)
@@ -398,7 +402,8 @@ class TestMetricTensor:
 
         layer2_diag = qml.QNode(layer2_diag, dev)
 
-        def layer2_off_diag_first_order(x, y, z, h, g, f):
+        def layer2_off_diag_first_order(params):
+            x, y, z, h, g, f = params
             non_parametrized_layer(a, b, c)
             qml.RX(x, wires=0)
             qml.RY(y, wires=1)
@@ -408,7 +413,8 @@ class TestMetricTensor:
 
         layer2_off_diag_first_order = qml.QNode(layer2_off_diag_first_order, dev)
 
-        def layer2_off_diag_second_order(x, y, z, h, g, f):
+        def layer2_off_diag_second_order(params):
+            x, y, z, h, g, f = params
             non_parametrized_layer(a, b, c)
             qml.RX(x, wires=0)
             qml.RY(y, wires=1)
@@ -419,13 +425,13 @@ class TestMetricTensor:
         layer2_off_diag_second_order = qml.QNode(layer2_off_diag_second_order, dev)
 
         # calculate the diagonal terms
-        varK0, varK1 = layer2_diag(x, y, z, h, g, f)
+        varK0, varK1 = layer2_diag(params)
         G2[0, 0] = varK0 / 4
         G2[1, 1] = varK1 / 4
 
         # calculate the off-diagonal terms
-        exK0, exK1 = layer2_off_diag_first_order(x, y, z, h, g, f)
-        exK01 = layer2_off_diag_second_order(x, y, z, h, g, f)
+        exK0, exK1 = layer2_off_diag_first_order(params)
+        exK01 = layer2_off_diag_second_order(params)
 
         G2[0, 1] = (exK01 - exK0 * exK1) / 4
         G2[1, 0] = (exK01 - exK0 * exK1) / 4
@@ -446,7 +452,8 @@ class TestMetricTensor:
         # Note: since this layer only consists of a single parameter,
         # only need to compute a single diagonal element.
 
-        def layer3_diag(x, y, z, h, g, f):
+        def layer3_diag(params):
+            x, y, z, h, g, f = params
             non_parametrized_layer(a, b, c)
             qml.RX(x, wires=0)
             qml.RY(y, wires=1)
@@ -456,7 +463,7 @@ class TestMetricTensor:
             return qml.var(qml.PauliX(1))
 
         layer3_diag = qml.QNode(layer3_diag, dev)
-        G3 = layer3_diag(x, y, z, h, g, f) / 4
+        G3 = layer3_diag(params) / 4
         assert qml.math.allclose(G[3:4, 3:4], G3, atol=tol, rtol=0)
 
         # ============================================
@@ -470,10 +477,12 @@ class TestMetricTensor:
         """Test that a metric tensor under the diagonal approximation evaluates
         correctly."""
         dev, circuit, non_parametrized_layer, a, b, c = sample_circuit
-        params = [-0.282203, 0.145554, 0.331624, -0.163907, 0.57662, 0.081272]
-        x, y, z, h, g, f = params
+        params = np.array(
+            [-0.282203, 0.145554, 0.331624, -0.163907, 0.57662, 0.081272],
+            requires_grad=True,
+        )
 
-        G = qml.metric_tensor(circuit, approx="diag")(*params)
+        G = qml.metric_tensor(circuit, approx="diag")(params)
 
         # ============================================
         # Test block-diag metric tensor of first layer is correct.
@@ -512,7 +521,8 @@ class TestMetricTensor:
         #   qml.RZ(g, wires=2)
         G2 = np.zeros([2, 2])
 
-        def layer2_diag(x, y, z, h, g, f):
+        def layer2_diag(params):
+            x, y, z, h, g, f = params
             non_parametrized_layer(a, b, c)
             qml.RX(x, wires=0)
             qml.RY(y, wires=1)
@@ -523,7 +533,7 @@ class TestMetricTensor:
         layer2_diag = qml.QNode(layer2_diag, dev)
 
         # calculate the diagonal terms
-        varK0, varK1 = layer2_diag(x, y, z, h, g, f)
+        varK0, varK1 = layer2_diag(params)
         G2[0, 0] = varK0 / 4
         G2[1, 1] = varK1 / 4
 
@@ -543,7 +553,8 @@ class TestMetricTensor:
         # Note: since this layer only consists of a single parameter,
         # only need to compute a single diagonal element.
 
-        def layer3_diag(x, y, z, h, g, f):
+        def layer3_diag(params):
+            x, y, z, h, g, f = params
             non_parametrized_layer(a, b, c)
             qml.RX(x, wires=0)
             qml.RY(y, wires=1)
@@ -553,7 +564,7 @@ class TestMetricTensor:
             return qml.var(qml.PauliX(1))
 
         layer3_diag = qml.QNode(layer3_diag, dev)
-        G3 = layer3_diag(x, y, z, h, g, f) / 4
+        G3 = layer3_diag(params) / 4
         assert qml.math.allclose(G[3:4, 3:4], G3, atol=tol, rtol=0)
 
         # ============================================
@@ -715,48 +726,38 @@ fubini_ansatze = [
     fubini_ansatz8,
 ]
 
+B = np.array(
+    [
+        [
+            [0.73, 0.49, 0.04],
+            [0.29, 0.45, 0.59],
+            [0.64, 0.06, 0.26],
+        ],
+        [
+            [0.93, 0.14, 0.46],
+            [0.31, 0.83, 0.79],
+            [0.25, 0.40, 0.16],
+        ],
+    ],
+    requires_grad=True,
+)
 fubini_params = [
     (np.array([0.3434, -0.7245345], requires_grad=True),),
-    (
-        np.reshape(
-            [
-                0.73,
-                0.49,
-                0.04,
-                0.29,
-                0.45,
-                0.59,
-                0.64,
-                0.06,
-                0.26,
-                0.93,
-                0.14,
-                0.46,
-                0.31,
-                0.83,
-                0.79,
-                0.25,
-                0.40,
-                0.16,
-            ],
-            (2, 3, 3),
-            requires_grad=True,
-        ),
-    ),
+    (B,),
     (np.array([-0.1111, -0.2222], requires_grad=True),),
     (np.array([-0.1111, -0.2222, 0.4554], requires_grad=True),),
     (
-        -0.1735,
+        np.array(-0.1735, requires_grad=True),
         np.array([-0.1735, -0.2846, -0.2846], requires_grad=True),
     ),
     (np.array([-0.1735, -0.2846], requires_grad=True),),
     (np.array([-0.1735, -0.2846], requires_grad=True),),
-    (-0.1735,),
+    (np.array(-0.1735, requires_grad=True),),
     (np.array([-0.1111, 0.3333], requires_grad=True),),
 ]
 
 
-def autodiff_metric_tensor(ansatz, num_wires):
+def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
     """Compute the metric tensor by full state vector
     differentiation via autograd."""
     dev = qml.device("default.qubit", wires=num_wires)
@@ -772,6 +773,14 @@ def autodiff_metric_tensor(ansatz, num_wires):
         iqnode = lambda *params: np.imag(qnode(*params))
         rjac = qml.jacobian(rqnode)(*params)
         ijac = qml.jacobian(iqnode)(*params)
+
+        if not mimic_autograd:
+            if len(params) > 1 and all(
+                qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:]
+            ):
+                rjac = tuple(qml.math.transpose(rjac))
+                ijac = tuple(qml.math.transpose(ijac))
+
         if isinstance(rjac, tuple):
             out = []
             for rc, ic in zip(rjac, ijac):
@@ -785,10 +794,18 @@ def autodiff_metric_tensor(ansatz, num_wires):
                 )
             return tuple(out)
 
+        if (
+            mimic_autograd
+            and len(params) > 1
+            and all(qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:])
+        ):
+            jac_contract = [-2]
+        else:
+            jac_contract = [0]
         jac = rjac + 1j * ijac
-        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], [0]))
+        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], jac_contract))
         return np.real(
-            np.tensordot(np.conj(jac), jac, axes=([0], [0]))
+            np.tensordot(np.conj(jac), jac, axes=(jac_contract, jac_contract))
             - np.tensordot(np.conj(psidpsi), psidpsi, axes=0)
         )
 
@@ -801,7 +818,7 @@ class TestFullMetricTensor:
 
     @pytest.mark.parametrize("ansatz, params", zip(fubini_ansatze, fubini_params))
     def test_correct_output_autograd(self, ansatz, params):
-        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        expected = autodiff_metric_tensor(ansatz, self.num_wires, True)(*params)
         dev = qml.device("default.qubit.autograd", wires=self.num_wires + 1)
 
         @qml.qnode(dev, interface="autograd")
@@ -1003,9 +1020,7 @@ class TestDifferentiability:
         def cost_full(*weights):
             return np.array(qml.metric_tensor(qnode, approx=None)(*weights))
 
-        _cost_full = lambda *weights: np.array(
-            autodiff_metric_tensor(ansatz, num_wires=3)(*weights)
-        )
+        _cost_full = lambda *weights: np.array(autodiff_metric_tensor(ansatz, 3, True)(*weights))
         _c = _cost_full(*weights)
         c = cost_full(*weights)
         assert all(
@@ -1254,7 +1269,7 @@ def test_error_generator_not_registered(allow_nonunitary, monkeypatch):
                 qml.metric_tensor(circuit, approx=None, allow_nonunitary=allow_nonunitary)(x, z)
 
 
-def test_no_error_missing_aux_wire_not_used():
+def test_no_error_missing_aux_wire_not_used(recwarn):
     """Tests that a no error is raised if the requested (or default, if not given)
     auxiliary wire for the Hadamard test is missing but it is not used, either
     because ``approx`` is used or because there only is a diagonal contribution."""
@@ -1289,6 +1304,8 @@ def test_no_error_missing_aux_wire_not_used():
     qml.metric_tensor(circuit_multi_block, approx="block-diag")(x, z)
     qml.metric_tensor(circuit_multi_block, approx="block-diag", aux_wire="aux_wire")(x, z)
 
+    assert len(recwarn) == 0
+
 
 def aux_wire_ansatz_0(x, y):
     qml.RX(x, wires=0)