Repository: PennyLaneAI/pennylane
Pull request #1725 (2021-10-07 16:22:50)

Add Hadamard test-based metric tensor
**Context:**
The metric tensor can be computed in a QPU-compatible way using Hadamard tests (see Fig. 3 in [1701.02691](https://arxiv.org/abs/1701.02691), Fig. 2 in [PhysRevX.7.021050](https://journals.aps.org/prx/pdf/10.1103/PhysRevX.7.021050) and Fig. 5 in [1804.03023](https://arxiv.org/pdf/1804.03023.pdf) ).
This allows for the computation of the full metric tensor in a hardware-ready manner.

**Description of the Change:**
The `metric_tensor` transform is extended so that users can compute the full tensor.
They may choose a block diagonal or diagonal approximation instead, reverting back to the current implementation of the tensor.
Even if the full tensor is requested, the block diagonal is computed via the old implementation, as it is cheaper and hardware compatible as well.

**Benefits:**
Finally the full metric tensor is available in PennyLane, in a QPU-ready and differentiable manner

**Possible Drawbacks:**
This changes the default behaviour of `metric_tensor` to compute the full metric now.
This is a breaking change but previously the approximation was rather silent, so that this is not really a drawback
except for compatibility issues and increased cost of old code as compared to the previous default behaviour.
The old default can be restored by using `approx="block-diag"`.

**Related GitHub Issues:**
The restructuring of the `metric_tensor` transform prepared this PR #1721 
The adjoint method for the metric_tensor transform also is WIP. It is cheaper on simulators but not available for QPUs and (currently) does not allow for shots != None. #1387 


josh146: [sc-9761]
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1725](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (2a7b0bb) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/d705f0e3cc51bd876c8b910cd25b85bfc08694b0?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (d705f0e) will **decrease** coverage by `0.02%`.
> The diff coverage is `94.33%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725/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/1725?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1725      +/-   ##
==========================================
- Coverage   98.84%   98.81%   -0.03%     
==========================================
  Files         224      224              
  Lines       17101    17199      +98     
==========================================
+ Hits        16903    16995      +92     
- Misses        198      204       +6     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/math/quantum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvcXVhbnR1bS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/transforms/metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWV0cmljX3RlbnNvci5weQ==) | `96.75% <94.33%> (-3.25%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725?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/1725?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [d705f0e...2a7b0bb](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1725?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).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 84cdcba56..b05755b4b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,55 @@
 
 <h3>New features since last release</h3>
 
+* The `metric_tensor` transform can now be used to compute the full
+  tensor, beyond the block diagonal approximation. 
+  [(#1725)](https://github.com/PennyLaneAI/pennylane/pull/1725)
+
+  This is performed using Hadamard tests, and requires an additional wire 
+  on the device to execute the circuits produced by the transform, 
+  as compared to the number of wires required by the original circuit.
+  The transform defaults to computing the full tensor, which can
+  be controlled by the `approx` keyword argument.
+  See the 
+  [qml.metric_tensor docstring](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.transforms.metric_tensor.html).
+  for more information and usage details.
+
+  As an example, consider the QNode
+
+  ```python
+  dev = qml.device("default.qubit", wires=3)
+
+  @qml.qnode(dev)
+  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)
+      return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+  weights = np.array([0.2, 1.2, -0.9], requires_grad=True)
+  ```
+
+  Then we can compute the (block) diagonal metric tensor as before, now using the
+  ``approx="block-diag"`` keyword:
+
+  ```pycon
+  >>> qml.metric_tensor(circuit, approx="block-diag")(weights)
+  [[0.25       0.         0.        ]
+   [0.         0.24013262 0.        ]
+   [0.         0.         0.21846983]]
+  ```
+
+  Instead, we now can also compute the full metric tensor, using
+  Hadamard tests on the additional wire of the device:
+
+  ```pycon
+  >>> qml.metric_tensor(circuit)(weights)
+  [[ 0.25        0.         -0.23300977]
+   [ 0.          0.24013262  0.01763859]
+   [-0.23300977  0.01763859  0.21846983]]
+  ```
+
 * Custom decompositions can now be applied to operations at the device level.
   [(#1900)](https://github.com/PennyLaneAI/pennylane/pull/1900)
 
@@ -354,6 +403,19 @@
 
 <h3>Breaking changes</h3>
 
+* The default behaviour of the `qml.metric_tensor` transform has been modified:
+  By default, the full metric tensor is computed, leading to higher cost than the previous
+  default of computing the block diagonal only. At the same time, the Hadamard tests for
+  the full metric tensor require an additional wire on the device, so that 
+
+  ```pycon
+  >>> qml.metric_tensor(some_qnode)(weights)
+  ```
+
+  will revert back to the block diagonal restriction and raise a warning if the
+  used device does not have an additional wire.
+  [(#1725)](https://github.com/PennyLaneAI/pennylane/pull/1725)
+
 * The `circuit_drawer` module has been renamed `drawer`.
   [(#1949)](https://github.com/PennyLaneAI/pennylane/pull/1949)
 
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index 558e42f02..883f09ba2 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -44,6 +44,7 @@ from .multi_dispatch import (
     ones_like,
     scatter_element_add,
     stack,
+    tensordot,
     unwrap,
     where,
 )
@@ -74,26 +75,27 @@ def __getattr__(name):
 
 __all__ = [
     "_multi_dispatch",
-    "block_diag",
-    "concatenate",
-    "diag",
-    "dot",
-    "ones_like",
-    "stack",
-    "where",
     "allclose",
     "allequal",
+    "block_diag",
     "cast",
     "cast_like",
-    "is_abstract",
+    "concatenate",
     "convert_like",
-    "get_interface",
-    "requires_grad",
     "cov_matrix",
-    "marginal_prob",
-    "scatter_element_add",
-    "unwrap",
+    "diag",
+    "dot",
     "frobenius_inner_product",
+    "get_interface",
     "get_trainable_indices",
+    "is_abstract",
     "is_independent",
+    "marginal_prob",
+    "ones_like",
+    "requires_grad",
+    "scatter_element_add",
+    "stack",
+    "tensordot",
+    "unwrap",
+    "where",
 ]
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index 5723c46e3..ae3461952 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -252,6 +252,36 @@ def dot(tensor1, tensor2):
     return np.dot(x, y, like=interface)
 
 
+def tensordot(tensor1, tensor2, axes=None):
+    """Returns the tensor product of two tensors.
+    In general ``axes`` specifies either the set of axes for both
+    tensors that are contracted (with the first/second entry of ``axes``
+    giving all axis indices for the first/second tensor) or --- if it is
+    an integer --- the number of last/first axes of the first/second
+    tensor to contract over.
+    There are some non-obvious special cases:
+
+    * If both tensors are 0-dimensional, ``axes`` must be 0.
+      and a 0-dimensional scalar is returned containing the simple product.
+
+    * If both tensors are 1-dimensional and ``axes=0``, the outer product
+      is returned.
+
+    * Products between a non-0-dimensional and a 0-dimensional tensor are not
+      supported in all interfaces.
+
+    Args:
+        tensor1 (tensor_like): input tensor
+        tensor2 (tensor_like): input tensor
+        axes (int or list[list[int]]): Axes to contract over, see detail description.
+
+    Returns:
+        tensor_like: the tensor product of the two input tensors
+    """
+    interface = _multi_dispatch([tensor1, tensor2])
+    return np.tensordot(tensor1, tensor2, axes=axes, like=interface)
+
+
 def get_trainable_indices(values):
     """Returns a set containing the trainable indices of a sequence of
     values.
diff --git a/pennylane/math/quantum.py b/pennylane/math/quantum.py
index 2bc65bfe0..c926638da 100644
--- a/pennylane/math/quantum.py
+++ b/pennylane/math/quantum.py
@@ -34,7 +34,7 @@ def cov_matrix(prob, obs, wires=None, diag_approx=False):
     Args:
         prob (tensor_like): probability distribution
         obs (list[.Observable]): a list of observables for which
-            to compute the covariance matrix for
+            to compute the covariance matrix
         diag_approx (bool): if True, return the diagonal approximation
         wires (.Wires): The wire register of the system. If not provided,
             it is assumed that the wires are labelled with consecutive integers.
diff --git a/pennylane/templates/subroutines/commuting_evolution.py b/pennylane/templates/subroutines/commuting_evolution.py
index 5221e0679..f0d46dc89 100644
--- a/pennylane/templates/subroutines/commuting_evolution.py
+++ b/pennylane/templates/subroutines/commuting_evolution.py
@@ -14,7 +14,7 @@
 r"""
 Contains the CommutingEvolution template.
 """
-# pylint: disable-msg=too-many-arguments
+# pylint: disable-msg=too-many-arguments,import-outside-toplevel
 import pennylane as qml
 from pennylane.operation import Operation, AnyWires
 
@@ -108,6 +108,8 @@ class CommutingEvolution(Operation):
     grad_method = None
 
     def __init__(self, hamiltonian, time, frequencies=None, shifts=None, do_queue=True, id=None):
+        from pennylane.gradients.general_shift_rules import get_shift_rule
+
         # pylint: disable=import-outside-toplevel
         from pennylane.gradients.general_shift_rules import (
             get_shift_rule,
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index a561b91f8..bf344e092 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -18,23 +18,47 @@ methods of computing the metric tensor.
 import functools
 import warnings
 
-import numpy as np
 import pennylane as qml
+from pennylane import numpy as np
+from pennylane.circuit_graph import LayerData
 
 from .batch_transform import batch_transform
 
-
-def expand_fn(tape, approx="block-diag", diag_approx=None, allow_nonunitary=True):
+# This dictionary maps generators to their controlled gate equivalent.
+# It is equivalent to using ``controlled`` on the generators, but this
+# approach via registering specific generators allows us to guarantee
+# unitary circuits for ``allow_nonunitary=False``
+_GEN_TO_CGEN = {
+    qml.PauliX: qml.CNOT,
+    qml.PauliY: qml.CY,
+    qml.PauliZ: qml.CZ,
+}
+
+# In contrast to the above, this dictionary maps *operators* to controlled
+# generator gates, skipping over the step of retrieving the generator from
+# the original operation and identifying it as a gate. It replaces matrix
+# arithmetics that otherwise would be required at this point.
+_OP_TO_CGEN = {
+    # PhaseShift is the same as RZ up to a global phase
+    qml.PhaseShift: qml.CZ,
+}
+
+
+def expand_fn(
+    tape, approx=None, diag_approx=None, allow_nonunitary=True, aux_wire=None, device_wires=None
+):
     """Set the metric tensor based on whether non-unitary gates are allowed."""
-    # pylint: disable=unused-argument
+    # pylint: disable=unused-argument,too-many-arguments
     if not allow_nonunitary and approx is None:  # pragma: no cover
         return qml.transforms.expand_nonunitary_gen(tape)
     return qml.transforms.expand_multipar(tape)
 
 
 @functools.partial(batch_transform, expand_fn=expand_fn)
-def metric_tensor(tape, approx="block-diag", diag_approx=None, allow_nonunitary=True):
-    """Returns a function that computes the block-diagonal approximation of the metric tensor
+def metric_tensor(
+    tape, approx=None, diag_approx=None, allow_nonunitary=True, aux_wire=None, device_wires=None
+):
+    r"""Returns a function that computes the block-diagonal approximation of the metric tensor
     of a given QNode or quantum tape.
 
     .. note::
@@ -42,44 +66,77 @@ def metric_tensor(tape, approx="block-diag", diag_approx=None, allow_nonunitary=
         Only gates that have a single parameter and define a ``generator`` are supported.
         All other parametrized gates will be decomposed if possible.
 
-    .. warning::
-
-        While ``approx=None`` is a valid input, the full metric tensor is not implemented yet
-        but will be added in an upcoming enhancement. Effectively, this means that only
-        ``approx="block-diag"`` and ``approx="diag"`` are currently supported.
-
     Args:
         tape (pennylane.QNode or .QuantumTape): quantum tape or QNode to find the metric tensor of
         approx (str): Which approximation of the metric tensor to compute.
 
             - If ``None``, the full metric tensor is computed
 
-            - If ``"block-diag"``, the block diagonal approximation is computed, reducing
+            - If ``"block-diag"``, the block-diagonal approximation is computed, reducing
               the number of evaluated circuits significantly.
 
             - If ``"diag"``, only the diagonal approximation is computed, slightly
-              reducing the classical overhead but not the quantum resources.
+              reducing the classical overhead but not the quantum resources
+              (compared to ``"block-diag"``).
 
-        diag_approx (bool): if True, use the diagonal approximation. If ``False``, a
-            block diagonal approximation of the metric tensor is computed.
-            This keyword argument is deprecated in favor of ``approx`` and will be removed soon
+        diag_approx (bool): if True, use the diagonal approximation. If ``False`` , a
+            block-diagonal approximation of the metric tensor is computed.
+            This keyword argument is deprecated in favor of ``approx`` and will be removed soon.
         allow_nonunitary (bool): Whether non-unitary operations are allowed in circuits
-            created by the transform. Only relevant if ``approx`` is ``None``
+            created by the transform. Only relevant if ``approx`` is ``None``.
+            Should be set to ``True`` if possible to reduce cost.
+        aux_wire (int or str or pennylane.wires.Wires): Auxiliary wire to be used for
+            Hadamard tests. If ``None`` (the default), a suitable wire is inferred
+            from the (number of) used wires in the original circuit and ``device_wires``.
+        device_wires (.wires.Wires): Wires of the device that is going to be used for the
+            metric tensor. Facilitates finding a default for ``aux_wire`` if ``aux_wire``
+            is ``None``.
         hybrid (bool): Specifies whether classical processing inside a QNode
             should be taken into account when transforming a QNode.
 
             - If ``True``, and classical processing is detected, the Jacobian of the classical
               processing will be computed and included. When evaluated, the
               returned metric tensor will be with respect to the QNode arguments.
+              The output shape can vary widely.
 
             - If ``False``, any internal QNode classical processing will be
               **ignored**. When evaluated, the returned metric tensor will be with
               respect to the **gate** arguments, and not the QNode arguments.
+              The output shape is a single two-dimensional tensor.
 
     Returns:
         func: Function which accepts the same arguments as the QNode. When called, this
         function will return the metric tensor.
 
+    The block-diagonal part of the metric tensor always is computed using the
+    covariance-based approach. If no approximation is selected,
+    the off block-diagonal is computed using Hadamard tests.
+
+    .. warning::
+
+        Performing the Hadamard tests requires a device
+        that has an additional wire as compared to the wires on which the
+        original circuit was defined. This wire may be specified via ``aux_wire``.
+        The available wires on the device may be specified via ``device_wires``.
+
+        By default (that is, if ``device_wires=None`` ), contiguous wire
+        numbering and usage is assumed and the additional
+        wire is set to the next wire of the device after the circuit wires.
+
+        If the given or inferred ``aux_wire`` does not exist on the device,
+        a warning is raised and the block-diagonal approximation is computed instead.
+        It is significantly cheaper in this case to explicitly set ``approx="block-diag"`` .
+
+    The flag ``allow_nonunitary`` should be set to ``True`` whenever the device with
+    which the metric tensor is computed supports non-unitary operations.
+    This will avoid additional decompositions of gates, in turn avoiding a potentially
+    large number of additional Hadamard test circuits to be run.
+    State vector simulators, for example, often allow applying operations that are
+    not unitary.
+    On a real QPU, setting this flag to ``True`` may cause exceptions because the
+    computation of the metric tensor will request invalid operations on a quantum
+    device.
+
     **Example**
 
     Consider the following QNode:
@@ -90,40 +147,50 @@ def metric_tensor(tape, approx="block-diag", diag_approx=None, allow_nonunitary=
 
         @qml.qnode(dev, interface="autograd")
         def circuit(weights):
-            # layer 1
             qml.RX(weights[0], wires=0)
-            qml.RX(weights[1], wires=1)
-
-            qml.CNOT(wires=[0, 1])
-            qml.CNOT(wires=[1, 2])
-
-            # layer 2
-            qml.RZ(weights[2], wires=0)
-            qml.RZ(weights[3], wires=2)
-
+            qml.RY(weights[1], wires=0)
             qml.CNOT(wires=[0, 1])
-            qml.CNOT(wires=[1, 2])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)), qml.expval(qml.PauliY(2))
+            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 ``metric_tensor`` transform to generate a new function that returns the
     metric tensor of this QNode:
 
-    >>> met_fn = qml.metric_tensor(circuit)
+    >>> mt_fn = qml.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)
+
+    In order to save cost, one might want to compute only the block-diagonal part of
+    the metric tensor, which requires significantly fewer executions of quantum functions
+    and does not need an auxiliary wire on the device. This can be done using the
+    ``approx`` keyword:
+
+    >>> mt_fn = qml.metric_tensor(circuit, approx="block-diag")
     >>> weights = np.array([0.1, 0.2, 0.4, 0.5], requires_grad=True)
-    >>> met_fn(weights)
+    >>> mt_fn(weights)
     tensor([[0.25  , 0.    , 0.    , 0.    ],
-            [0.    , 0.25  , 0.    , 0.    ],
-            [0.    , 0.    , 0.0025, 0.0024],
-            [0.    , 0.    , 0.0024, 0.0123]], requires_grad=True)
+            [0.    , 0.2475, 0.    , 0.    ],
+            [0.    , 0.    , 0.0123, 0.0123],
+            [0.    , 0.    , 0.0123, 0.0123]], requires_grad=True)
+
+    These blocks are given by parameter groups that belong to groups of commuting gates.
+
+    The tensor can be further restricted to the diagonal via ``approx="diag"``. However,
+    this will not save further quantum function evolutions but only classical postprocessing.
 
     The returned metric tensor is also fully differentiable in all interfaces.
-    For example, we can compute the gradient of the ``(3, 2)`` element
-    with respect to the QNode ``weights``:
+    For example, we can compute the gradient of the Frobenius norm of the metric tensor
+    with respect to the QNode ``weights`` :
 
-    >>> grad_fn = qml.grad(lambda x: met_fn(x)[3, 2])
+    >>> norm_fn = lambda x: qml.math.linalg(mt_fn(x), ord="fro")
+    >>> grad_fn = qml.grad(norm_fn)
     >>> grad_fn(weights)
-    array([[ 0.04867729, -0.00049502,  0.        ],
-           [ 0.        ,  0.        ,  0.        ]])
+    array([-0.0282246 ,  0.01340413,  0.        ,  0.        ])
 
     .. UsageDetails::
 
@@ -143,21 +210,45 @@ def metric_tensor(tape, approx="block-diag", diag_approx=None, allow_nonunitary=
         >>> tapes
         [<QuantumTape: wires=[0, 1], params=0>,
          <QuantumTape: wires=[0, 1], params=1>,
-         <QuantumTape: wires=[0, 1], params=3>]
+         <QuantumTape: wires=[0, 1], params=3>,
+         <QuantumTape: wires=[2, 0], params=1>,
+         <QuantumTape: wires=[2, 0, 1], params=2>,
+         <QuantumTape: wires=[2, 0, 1], params=2>]
 
         This can be useful if the underlying circuits representing the metric tensor
-        computation need to be analyzed.
+        computation need to be analyzed. We clearly can distinguish the first three
+        tapes used for the block-diagonal from the last three tapes that use the
+        auxiliary wire ``2`` , which was not used by the original tape.
 
         The output tapes can then be evaluated and post-processed to retrieve
         the metric tensor:
 
         >>> dev = qml.device("default.qubit", wires=2)
         >>> fn(qml.execute(tapes, dev, None))
-        array([[0.25      , 0.        , 0.        ],
-               [0.        , 0.00415023, 0.        ],
-               [0.        , 0.        , 0.24878844]])
+        array([[ 0.25      ,  0.        ,  0.42073549],
+               [ 0.        ,  0.00415023, -0.26517488],
+               [ 0.42073549, -0.26517488,  0.24878844]])
+
+        The first term of the off block-diagonal entries of the full metric tensor are
+        computed with Hadamard tests. This first term reads
+
+        .. math ::
+
+            \mathfrak{Re}\left\{\langle \partial_i\psi|\partial_j\psi\rangle\right\}
+
+        and can be computed using an augmented circuit with an additional qubit.
+        See for example `the appendix of McArdle et al. <https://arxiv.org/pdf/1804.03023.pdf>`__
+        for details.
+        The block-diagonal of the tensor is computed using the covariance matrix approach.
+
+        In addition, we may extract the factors for the second terms
+        :math:`\langle \psi|\partial_j\psi\rangle`
+        of the *off block-diagonal* tensor from the quantum function output for the covariance matrix!
+
+        This means that in total only the tapes for the first terms of the off block-diagonal
+        are required in addition to the circuits for the block diagonal.
     """
-    # pylint: disable=unused-argument
+    # pylint: disable=too-many-arguments
     if diag_approx is not None:
         warnings.warn(
             "The keyword argument diag_approx is deprecated. Please use approx='diag' instead.",
@@ -169,9 +260,15 @@ def metric_tensor(tape, approx="block-diag", diag_approx=None, allow_nonunitary=
     if approx in {"diag", "block-diag"}:
         # Only require covariance matrix based transform
         diag_approx = approx == "diag"
-        return _metric_tensor_cov_matrix(tape, diag_approx)
+        return _metric_tensor_cov_matrix(tape, diag_approx)[:2]
 
-    raise NotImplementedError("No method for the full metric tensor has been implemented yet.")
+    if approx is None:
+        return _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires)
+
+    raise ValueError(
+        f"Unknown value {approx} for keyword argument approx. "
+        "Valid values are 'diag', 'block-diag' and None."
+    )
 
 
 @metric_tensor.custom_qnode_wrapper
@@ -191,13 +288,34 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
 
         qnode = qnode.qnodes.qnodes[0]
 
+    tkwargs.setdefault("device_wires", qnode.device.wires)
     mt_fn = self.default_qnode_wrapper(qnode, targs, tkwargs)
 
     _expand_fn = lambda tape: self.expand_fn(tape, *targs, **tkwargs)
     cjac_fn = qml.transforms.classical_jacobian(qnode, expand_fn=_expand_fn)
 
     def wrapper(*args, **kwargs):
-        mt = mt_fn(*args, **kwargs)
+        try:
+            mt = mt_fn(*args, **kwargs)
+        except qml.wires.WireError as e:
+            if str(e) == "No device wires are unused by the tape.":
+                warnings.warn(
+                    "The device does not have a wire that is not used by the tape."
+                    "\n\nReverting to the block-diagonal approximation. It will often be "
+                    "much more efficient to request the block-diagonal approximation directly!"
+                )
+            else:
+                warnings.warn(
+                    "An auxiliary wire is not available."
+                    "\n\nThis can occur when computing the full metric tensor via the "
+                    "Hadamard test, and the device does not provide an "
+                    "additional wire or the requested auxiliary wire does not exist "
+                    "on the device."
+                    "\n\nReverting to the block-diagonal approximation. It will often be "
+                    "much more efficient to request the block-diagonal approximation directly!"
+                )
+            tkwargs["approx"] = "block-diag"
+            return self.__call__(qnode, *targs, **tkwargs)(*args, **kwargs)
 
         if not hybrid:
             return mt
@@ -209,13 +327,13 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
             if len(cjac) == 1:
                 cjac = cjac[0]
             else:
-                # Classical processing of multiple arguments is present. Return mt @ cjac.
+                # 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, [[-1], [0]])
-                        _mt = qml.math.tensordot(c, _mt, [[0], [0]])
+                        _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)
@@ -228,7 +346,6 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
             return mt
 
         # Classical processing of a single argument is present. Return mt @ cjac.
-        cjac = qml.math.convert_like(cjac, mt)
         mt = qml.math.tensordot(mt, cjac, [[-1], [0]])
         mt = qml.math.tensordot(cjac, mt, [[0], [0]])
         return mt
@@ -237,8 +354,23 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
 
 
 def _metric_tensor_cov_matrix(tape, diag_approx):
-    """This is the metric tensor method for the block diagonal, using
-    the covariance matrix of the generators of each layer."""
+    r"""This is the metric tensor method for the block diagonal, using
+    the covariance matrix of the generators of each layer.
+
+    Args:
+        tape (pennylane.QNode or .QuantumTape): quantum tape or QNode to find the metric tensor of
+        diag_approx (bool): if True, use the diagonal approximation. If ``False`` , a
+            block-diagonal approximation of the metric tensor is computed.
+    Returns:
+        list[pennylane.tape.QuantumTape]: Transformed tapes that compute the probabilities
+            required for the covariance matrix
+        callable: Post-processing function that computes the covariance matrix from the
+            results of the tapes in the first return value
+        list[list[.Observable]]: Observables measured in each tape, one inner list
+            corresponding to one tape in the first return value
+        list[list[float]]: Coefficients to scale the results for each observable, one inner list
+            corresponding to one tape in the first return value
+    """
     # get the circuit graph
     graph = tape.graph
 
@@ -269,8 +401,8 @@ def _metric_tensor_cov_matrix(tape, diag_approx):
 
             else:
                 raise qml.QuantumFunctionError(
-                    "Can't generate metric tensor, generator {}"
-                    "has no corresponding observable".format(gen)
+                    f"Can't generate metric tensor, generator {gen}"
+                    "has no corresponding observable"
                 )
 
         # Create a quantum tape with all operations
@@ -291,8 +423,9 @@ def _metric_tensor_cov_matrix(tape, diag_approx):
         gs = []
 
         for prob, obs, coeffs in zip(probs, obs_list, coeffs_list):
+
             # calculate the covariance matrix of this layer
-            scale = qml.math.convert_like(np.outer(coeffs, coeffs), prob)
+            scale = qml.math.convert_like(qml.math.outer(coeffs, coeffs), prob)
             scale = qml.math.cast_like(scale, prob)
             g = scale * qml.math.cov_matrix(prob, obs, wires=tape.wires, diag_approx=diag_approx)
             gs.append(g)
@@ -300,4 +433,250 @@ def _metric_tensor_cov_matrix(tape, diag_approx):
         # create the block diagonal metric tensor
         return qml.math.block_diag(gs)
 
-    return metric_tensor_tapes, processing_fn
+    return metric_tensor_tapes, processing_fn, obs_list, coeffs_list
+
+
+@functools.lru_cache()
+def _get_gen_op(op, allow_nonunitary, aux_wire):
+    r"""Get the controlled-generator operation for a given operation.
+
+    Args:
+        op (pennylane.operation.Operation): Operation from which to extract the generator
+        allow_nonunitary (bool): Whether non-unitary gates are allowed in the circuit
+        aux_wire (int or pennylane.wires.Wires): Auxiliary wire on which to control the operation
+
+    Returns
+        qml.Operation: Controlled-generator operation of the generator of ``op``, controlled
+        on wire ``aux_wire``.
+
+    Raises
+        ValueError: If the generator of ``op`` is not known or it is non-unitary while
+        ``allow_nonunitary=False``.
+
+    If ``allow_nonunitary=True``, a general :class:`~.pennylane.ControlledQubitUnitary` is returned,
+    otherwise only controlled Pauli operations are used. If the operation has a non-unitary
+    generator but ``allow_nonunitary=False``, the operation ``op`` should have been decomposed
+    before, leading to a ``ValueError``.
+    """
+    gen, _ = op.generator
+    try:
+        if isinstance(gen, np.ndarray):
+            cgen = _OP_TO_CGEN[op.__class__]
+        else:
+            cgen = _GEN_TO_CGEN.get(gen, None)
+            if cgen is None:
+                cgen = _OP_TO_CGEN[op.__class__]
+        return cgen(wires=[aux_wire, *op.wires])
+
+    except KeyError as e:
+        if allow_nonunitary:
+            if (not isinstance(gen, np.ndarray)) and issubclass(gen, qml.operation.Observable):
+                gen = gen.matrix
+            return qml.ControlledQubitUnitary(gen, control_wires=aux_wire, wires=op.wires)
+        raise ValueError(
+            f"Generator for operation {op} not known and non-unitary operations "
+            "deactivated via allow_nonunitary=False."
+        ) from e
+
+
+def _get_first_term_tapes(tape, layer_i, layer_j, allow_nonunitary, aux_wire):
+    r"""Obtain the tapes for the first term of all tensor entries
+    belonging to an off-diagonal block.
+
+    Args:
+        tape (pennylane.tape.QuantumTape): Tape that is being transformed
+        layer_i (list): The first layer of parametrized ops, of the format of
+            the layers generated by ``iterate_parametrized_layers``
+        layer_j (list): The second layer of parametrized ops
+        allow_nonunitary (bool): Whether non-unitary operations are allowed
+            in the circuit; passed to ``_get_gen_op``
+        aux_wire (object or pennylane.wires.Wires): Auxiliary wire on which to
+            control the controlled-generator operations
+
+    Returns:
+        list[pennylane.tape.QuantumTape]: Transformed tapes that compute the
+            first term of the metric tensor for the off-diagonal block belonging
+            to the input layers
+        list[tuple[int]]: 2-tuple indices assigning the tapes to metric tensor
+            entries
+    """
+
+    tapes = []
+    ids = []
+    # Exclude the backwards cone of layer_i from the backwards cone of layer_j
+    ops_between_cgens = [op for op in layer_j.pre_ops if op not in layer_i.pre_ops]
+
+    # Iterate over differentiated operation in first layer
+    for diffed_op_i, par_idx_i in zip(layer_i.ops, layer_i.param_inds):
+        gen_op_i = _get_gen_op(diffed_op_i, allow_nonunitary, aux_wire)
+
+        # Iterate over differentiated operation in second layer
+        # There will be one tape per pair of differentiated operations
+        for diffed_op_j, par_idx_j in zip(layer_j.ops, layer_j.param_inds):
+            gen_op_j = _get_gen_op(diffed_op_j, allow_nonunitary, aux_wire)
+
+            with tape.__class__() as new_tape:
+                # Initialize auxiliary wire
+                qml.Hadamard(wires=aux_wire)
+                # Apply backward cone of first layer
+                for op in layer_i.pre_ops:
+                    qml.apply(op)
+                # Controlled-generator operation of first diff'ed op
+                qml.apply(gen_op_i)
+                # Apply first layer and operations between layers
+                for op in ops_between_cgens:
+                    qml.apply(op)
+                # Controlled-generator operation of second diff'ed op
+                qml.apply(gen_op_j)
+                # Measure X on auxiliary wire
+                qml.expval(qml.PauliX(aux_wire))
+
+            tapes.append(new_tape)
+            # Memorize to which metric entry this tape belongs
+            ids.append((par_idx_i, par_idx_j))
+
+    return tapes, ids
+
+
+def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
+    r"""Generate the quantum tapes that execute the Hadamard tests
+    to compute the first term of off block-diagonal metric entries
+    and combine them with the covariance matrix-based block-diagonal tapes.
+
+    Args:
+        tape (pennylane.QNode or .QuantumTape): quantum tape or QNode to find the metric tensor of
+        allow_nonunitary (bool): Whether non-unitary operations are allowed in circuits
+            created by the transform. Only relevant if ``approx`` is ``None``
+            Should be set to ``True`` if possible to reduce cost.
+        aux_wire (int or .wires.Wires): Auxiliary wire to be used for
+            Hadamard tests. By default, a suitable wire is inferred from the number
+            of used wires in the original circuit.
+        device_wires (.wires.Wires): Wires of the device that is going to be used for the
+            metric tensor. Facilitates finding a default for ``aux_wire`` if ``aux_wire``
+            is ``None`` .
+
+    Returns:
+        list[pennylane.tape.QuantumTape]: Tapes to evaluate the metric tensor
+        callable: processing function to obtain the metric tensor from the tape results
+    """
+    # Get tapes and processing function for the block-diagonal metric tensor,
+    # as well as the generator observables and generator coefficients for each diff'ed operation
+    diag_tapes, diag_proc_fn, obs_list, coeffs = _metric_tensor_cov_matrix(tape, diag_approx=False)
+
+    # Obtain layers of parametrized operations and account for the discrepancy between trainable
+    # 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,
+            layer.ops,
+            tuple(par_idx_to_trainable_idx[idx] for idx in layer[2]),
+            layer.post_ops,
+        )
+        for layer in layers
+    ]
+
+    # Get default for aux_wire
+    aux_wire = _get_aux_wire(aux_wire, tape, device_wires)
+
+    # Get all tapes for the first term of the metric tensor and memorize which
+    # entry they belong to
+    first_term_tapes = []
+    ids = []
+    block_sizes = []
+    for idx_i, layer_i in enumerate(layers):
+        block_sizes.append(len(layer_i.param_inds))
+
+        for layer_j in layers[idx_i + 1 :]:
+            _tapes, _ids = _get_first_term_tapes(tape, layer_i, layer_j, allow_nonunitary, aux_wire)
+            first_term_tapes.extend(_tapes)
+            ids.extend(_ids)
+
+    # Combine block-diagonal and off block-diagonal tapes
+    tapes = diag_tapes + first_term_tapes
+    # prepare off block-diagonal mask
+    mask = 1 - qml.math.block_diag([qml.math.ones((bsize, bsize)) for bsize in block_sizes])
+    # Required for slicing in processing_fn
+    num_diag_tapes = len(diag_tapes)
+
+    def processing_fn(results):
+        """Postprocessing function for the full metric tensor."""
+        nonlocal mask
+        # Split results
+        diag_res, off_diag_res = results[:num_diag_tapes], results[num_diag_tapes:]
+        # Get full block-diagonal tensor
+        diag_mt = diag_proc_fn(diag_res)
+
+        # Prepare the mask to match the used interface
+        mask = qml.math.convert_like(mask, diag_mt)
+
+        # Initialize off block-diagonal tensor using the stored ids
+        first_term = qml.math.zeros_like(diag_mt)
+        if ids:
+            off_diag_res = qml.math.stack(off_diag_res, 1)[0]
+            inv_ids = [_id[::-1] for _id in ids]
+            first_term = qml.math.scatter_element_add(first_term, list(zip(*ids)), off_diag_res)
+            first_term = qml.math.scatter_element_add(first_term, list(zip(*inv_ids)), off_diag_res)
+
+        # Second terms of off block-diagonal metric tensor
+        expvals = qml.math.zeros_like(first_term[0])
+        idx = 0
+        for prob, obs in zip(diag_res, obs_list):
+            for o in obs:
+                l = qml.math.cast(o.eigvals, dtype=np.float64)
+                w = tape.wires.indices(o.wires)
+                p = qml.math.marginal_prob(prob, w)
+                expvals = qml.math.scatter_element_add(expvals, (idx,), qml.math.dot(l, p))
+                idx += 1
+
+        # Construct <\partial_i\psi|\psi><\psi|\partial_j\psi> and mask it
+        second_term = qml.math.tensordot(expvals, expvals, axes=0) * mask
+
+        # Subtract second term from first term
+        off_diag_mt = first_term - second_term
+
+        # Rescale first and second term
+        _coeffs = qml.math.hstack(coeffs)
+        scale = qml.math.convert_like(qml.math.tensordot(_coeffs, _coeffs, axes=0), results[0])
+        off_diag_mt = scale * off_diag_mt
+
+        # Combine block-diagonal and off block-diagonal
+        mt = off_diag_mt + diag_mt
+
+        return mt
+
+    return tapes, processing_fn
+
+
+def _get_aux_wire(aux_wire, tape, device_wires):
+    r"""Determine an unused wire to be used as auxiliary wire for Hadamard tests.
+
+    Args:
+        aux_wire (object): Input auxiliary wire. Returned unmodified if not ``None``
+        tape (pennylane.tape.QuantumTape): Tape to infer the wire for
+        device_wires (.wires.Wires): Wires of the device that is going to be used for the
+            metric tensor. Facilitates finding a default for ``aux_wire`` if ``aux_wire``
+            is ``None`` .
+
+    Returns:
+        object: The auxiliary wire to be used. Equals ``aux_wire`` if it was not ``None`` ,
+        and an often reasonable choice else.
+    """
+    if aux_wire is not None:
+        if device_wires is None or aux_wire in device_wires:
+            return aux_wire
+
+    if device_wires is not None:
+        unused_wires = qml.wires.Wires(device_wires.toset().difference(tape.wires.toset()))
+        if not unused_wires:
+            raise qml.wires.WireError("No device wires are unused by the tape.")
+        return unused_wires[0]
+
+    _wires = tape.wires
+    for _aux in range(tape.num_wires):
+        if _aux not in _wires:
+            return _aux
+
+    return tape.num_wires


Test patch:
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 37e9222ee..5e9d8c58b 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -535,8 +535,8 @@ class TestDot:
         assert fn.allequal(res, expected)
 
 
-class TestTensordot:
-    """Tests for the dot product function."""
+class TestTensordotTorch:
+    """Tests for the tensor product function in torch."""
 
     v1 = torch.tensor([0.1, 0.5, -0.9, 1.0, -4.2, 0.1], dtype=torch.float64)
     v2 = torch.tensor([4.3, -1.2, 8.2, 0.6, -4.2, -11.0], dtype=torch.float64)
@@ -558,7 +558,7 @@ class TestTensordot:
     M1 = torch.tensor(_arange)
     M2 = torch.tensor(_shuffled_arange)
     T1 = np.arange(0, 3 * 6 * 9 * 2).reshape((3, 6, 9, 2)).astype(np.float64)
-    T1 = np.array([T1[1], T1[0], T1[2]])
+    T1 = torch.tensor(np.array([T1[1], T1[0], T1[2]]), dtype=torch.float64)
 
     v1_dot_v2 = 9.59
     v1_outer_v2 = np.array(
@@ -1380,6 +1380,26 @@ class TestScatterElementAdd:
         assert fn.allclose(jac[0], self.expected_jac_x)
         assert fn.allclose(jac[1], self.expected_jac_y)
 
+    def test_array_batch(self):
+        """Test that a NumPy array and the addend are differentiable when using
+        scatter addition (multi dispatch)."""
+        x = np.ones((2, 3), requires_grad=True)
+        y = np.array([0.56, 0.3], requires_grad=True)
+
+        def cost_multi(weight_0, weight_1):
+            return fn.scatter_element_add(weight_0, [(0, 1), (1, 2)], weight_1 ** 2)
+
+        res = cost_multi(x, y)
+        assert isinstance(res, np.ndarray)
+        assert fn.allclose(res, onp.array([[1.0, 1.3136, 1.0], [1.0, 1.0, 1.09]]))
+
+        jac = qml.jacobian(lambda *weights: cost_multi(*weights))(x, y)
+        assert fn.allclose(jac[0], self.expected_jac_x)
+        exp_jac_y = onp.zeros((2, 3, 2))
+        exp_jac_y[0, 1, 0] = 2 * y[0]
+        exp_jac_y[1, 2, 1] = 2 * y[1]
+        assert fn.allclose(jac[1], exp_jac_y)
+
     def test_tensorflow(self):
         """Test that a TF tensor is differentiable when using scatter addition"""
         x = tf.Variable(self.x)
@@ -1534,8 +1554,6 @@ class TestScatterElementAddMultiValue:
 
         res = cost([x, y])
         assert isinstance(res, jax.interpreters.xla.DeviceArray)
-        print(res)
-        print(self.expected_val)
         assert fn.allclose(res, self.expected_val)
 
         scalar_cost = (
diff --git a/tests/optimize/test_qng.py b/tests/optimize/test_qng.py
index bfcd10035..bb1683140 100644
--- a/tests/optimize/test_qng.py
+++ b/tests/optimize/test_qng.py
@@ -179,15 +179,16 @@ class TestOptimize:
         # check final cost
         assert np.allclose(circuit(theta), -0.9963791, atol=tol, rtol=0)
 
-    def test_single_qubit_vqe(self, tol):
+    @pytest.mark.parametrize("approx", [None, "block-diag"])
+    def test_single_qubit_vqe(self, tol, approx):
         """Test single-qubit VQE has the correct QNG value
         every step, the correct parameter updates,
         and correct cost after 200 steps"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit", wires=(2 if approx is None else 1))
 
-        def circuit(params, wires=0):
-            qml.RX(params[0], wires=wires)
-            qml.RY(params[1], wires=wires)
+        def circuit(params, wires=None):
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=0)
 
         coeffs = [1, 1]
         obs_list = [qml.PauliX(0), qml.PauliZ(0)]
@@ -211,7 +212,7 @@ class TestOptimize:
         # optimization for 200 steps total
         for t in range(num_steps):
             theta_new = opt.step(
-                cost_fn, theta, metric_tensor_fn=qml.metric_tensor(qnodes.qnodes[0])
+                cost_fn, theta, metric_tensor_fn=qml.metric_tensor(qnodes.qnodes[0], approx=approx)
             )
 
             # check metric tensor
diff --git a/tests/test_vqe.py b/tests/test_vqe.py
index 4bb80fc19..2d731f7af 100644
--- a/tests/test_vqe.py
+++ b/tests/test_vqe.py
@@ -586,10 +586,11 @@ class TestVQE:
 
         assert np.allclose(dc, big_hamiltonian_grad)
 
-    def test_metric_tensor(self):
+    @pytest.mark.parametrize("approx", [None, "block-diag", "diag"])
+    def test_metric_tensor(self, approx):
         """Test that the metric tensor can be calculated."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit", wires=3)
         p = np.array([1.0, 1.0, 1.0])
 
         def ansatz(params, **kwargs):
@@ -600,7 +601,7 @@ class TestVQE:
 
         h = qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(1)])
         qnodes = qml.ExpvalCost(ansatz, h, dev)
-        mt = qml.metric_tensor(qnodes)(p)
+        mt = qml.metric_tensor(qnodes, approx=approx)(p)
         assert mt.shape == (3, 3)
         assert isinstance(mt, np.ndarray)
 
@@ -629,7 +630,7 @@ class TestVQE:
         assert np.allclose(res, exp)
 
         with pytest.warns(UserWarning, match="ExpvalCost was instantiated with multiple devices."):
-            qml.metric_tensor(qnodes)(w)
+            qml.metric_tensor(qnodes, approx="block-diag")(w)
 
     def test_multiple_devices_opt_true(self):
         """Test if a ValueError is raised when multiple devices are passed when optimize=True."""
diff --git a/tests/transforms/test_batch_transform.py b/tests/transforms/test_batch_transform.py
index 46d87a9f5..1fb9e4610 100644
--- a/tests/transforms/test_batch_transform.py
+++ b/tests/transforms/test_batch_transform.py
@@ -213,7 +213,6 @@ class TestBatchTransform:
         transform_fn = qml.batch_transform(
             MyTransform().my_transform, expand_fn=self.expand_logic_with_kwarg
         )
-
         with qml.tape.QuantumTape() as tape:
             qml.PhaseShift(0.5, wires=0)
             qml.expval(qml.PauliX(0))
@@ -260,7 +259,6 @@ class TestBatchTransform:
 
         spy_transform.assert_called()
         spy_expand.assert_called()  # The expand_fn of transform_fn always is called
-
         input_tape = spy_transform.call_args[0][1]
         assert len(input_tape.operations) == 1
         assert input_tape.operations[0].name == ("RZ" if perform_expansion else "PhaseShift")
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index 39b0c0f1b..1e9c0d71e 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -20,6 +20,7 @@ from scipy.linalg import block_diag
 
 import pennylane as qml
 from gate_data import Y, Z
+from pennylane.transforms.metric_tensor import _get_aux_wire
 
 
 class TestMetricTensor:
@@ -91,19 +92,6 @@ class TestMetricTensor:
         result = qml.metric_tensor(circuit, hybrid=False, approx="block-diag")(*params)
         assert result.shape == (2, 2)
 
-    def test_generator_no_expval(self, monkeypatch):
-        """Test exception is raised if subcircuit contains an
-        operation with generator object that is not an observable"""
-        with monkeypatch.context() as m:
-            m.setattr("pennylane.RX.generator", [qml.RX, 1])
-
-            with qml.tape.QuantumTape() as tape:
-                qml.RX(np.array(0.5, requires_grad=True), wires=0)
-                qml.expval(qml.PauliX(0))
-
-            with pytest.raises(qml.QuantumFunctionError, match="no corresponding observable"):
-                qml.metric_tensor(tape, approx="block-diag")
-
     def test_construct_subcircuit(self):
         """Test correct subcircuits constructed"""
         dev = qml.device("default.qubit", wires=2)
@@ -215,7 +203,8 @@ class TestMetricTensor:
         assert isinstance(tapes[3].operations[12], qml.Hadamard)
 
     def test_evaluate_diag_metric_tensor(self, tol):
-        """Test that a diagonal metric tensor evaluates correctly"""
+        """Test that a diagonal metric tensor evaluates correctly for
+        block-diagonal and diagonal setting."""
         dev = qml.device("default.qubit", wires=2)
 
         def circuit(a, b, c):
@@ -232,7 +221,8 @@ class TestMetricTensor:
         c = -0.432
 
         # evaluate metric tensor
-        g = qml.metric_tensor(circuit, approx="block-diag")(a, b, c)
+        g_diag = qml.metric_tensor(circuit, approx="diag")(a, b, c)
+        g_blockdiag = qml.metric_tensor(circuit, approx="block-diag")(a, b, c)
 
         # check that the metric tensor is correct
         expected = (
@@ -241,7 +231,8 @@ class TestMetricTensor:
             )
             / 4
         )
-        assert np.allclose(g, np.diag(expected), atol=tol, rtol=0)
+        assert qml.math.allclose(g_diag, np.diag(expected), atol=tol, rtol=0)
+        assert qml.math.allclose(g_blockdiag, np.diag(expected), atol=tol, rtol=0)
 
     @pytest.mark.parametrize("strategy", ["gradient", "device"])
     def test_template_integration(self, strategy, tol):
@@ -258,20 +249,6 @@ class TestMetricTensor:
         res = qml.metric_tensor(circuit, approx="block-diag")(weights)
         assert res.shape == (2, 3, 3, 2, 3, 3)
 
-    def test_full_tensor_not_implemented(self):
-        """Test that the full metric tensor can not be computed yet."""
-        dev = qml.device("default.qubit", wires=3)
-
-        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)
-        params = [0.1, 0.2]
-        with pytest.raises(NotImplementedError):
-            result = qml.metric_tensor(circuit, approx=None)(*params)
-
     def test_evaluate_diag_metric_tensor_classical_processing(self, tol):
         """Test that a diagonal metric tensor evaluates correctly
         when the QNode includes classical processing."""
@@ -302,10 +279,10 @@ class TestMetricTensor:
 
         # check that the metric tensor is correct
         expected = np.array([np.cos(a[1]) ** 2, 1]) / 4
-        assert np.allclose(g[0], np.diag(expected), atol=tol, rtol=0)
+        assert qml.math.allclose(g[0], np.diag(expected), atol=tol, rtol=0)
 
         expected = (3 - 2 * np.cos(a[1]) ** 2 * np.cos(2 * a[0]) - np.cos(2 * a[1])) / 16
-        assert np.allclose(g[1], expected, atol=tol, rtol=0)
+        assert qml.math.allclose(g[1], expected, atol=tol, rtol=0)
 
     @pytest.fixture(params=["parameter-shift", "backprop"])
     def sample_circuit(self, request):
@@ -397,7 +374,7 @@ class TestMetricTensor:
         G1[2, 0] = G1[0, 2]
         G1[2, 1] = G1[1, 2]
 
-        assert np.allclose(G[:3, :3], G1, atol=tol, rtol=0)
+        assert qml.math.allclose(G[:3, :3], G1, atol=tol, rtol=0)
 
         # =============================================
         # Test block-diag metric tensor of second layer is correct.
@@ -453,7 +430,7 @@ class TestMetricTensor:
         G2[0, 1] = (exK01 - exK0 * exK1) / 4
         G2[1, 0] = (exK01 - exK0 * exK1) / 4
 
-        assert np.allclose(G[4:6, 4:6], G2, atol=tol, rtol=0)
+        assert qml.math.allclose(G[4:6, 4:6], G2, atol=tol, rtol=0)
 
         # =============================================
         # Test block-diag metric tensor of third layer is correct.
@@ -480,14 +457,14 @@ class TestMetricTensor:
 
         layer3_diag = qml.QNode(layer3_diag, dev)
         G3 = layer3_diag(x, y, z, h, g, f) / 4
-        assert np.allclose(G[3:4, 3:4], G3, atol=tol, rtol=0)
+        assert qml.math.allclose(G[3:4, 3:4], G3, atol=tol, rtol=0)
 
         # ============================================
         # Finally, double check that the entire metric
         # tensor is as computed.
 
         G_expected = block_diag(G1, G3, G2)
-        assert np.allclose(G, G_expected, atol=tol, rtol=0)
+        assert qml.math.allclose(G, G_expected, atol=tol, rtol=0)
 
     def test_evaluate_diag_approx_metric_tensor(self, sample_circuit, tol):
         """Test that a metric tensor under the
@@ -524,8 +501,8 @@ class TestMetricTensor:
         ) / 64
         G1[2, 2] = (3 - np.cos(2 * a) - 2 * np.cos(a) ** 2 * np.cos(2 * (b + c))) / 16
 
-        assert np.allclose(G[:3, :3], G1, atol=tol, rtol=0)
-        assert np.allclose(G_alias[:3, :3], G1, atol=tol, rtol=0)
+        assert qml.math.allclose(G[:3, :3], G1, atol=tol, rtol=0)
+        assert qml.math.allclose(G_alias[:3, :3], G1, atol=tol, rtol=0)
 
         # =============================================
         # Test block-diag metric tensor of second layer is correct.
@@ -554,8 +531,8 @@ class TestMetricTensor:
         G2[0, 0] = varK0 / 4
         G2[1, 1] = varK1 / 4
 
-        assert np.allclose(G[4:6, 4:6], G2, atol=tol, rtol=0)
-        assert np.allclose(G_alias[4:6, 4:6], G2, atol=tol, rtol=0)
+        assert qml.math.allclose(G[4:6, 4:6], G2, atol=tol, rtol=0)
+        assert qml.math.allclose(G_alias[4:6, 4:6], G2, atol=tol, rtol=0)
 
         # =============================================
         # Test metric tensor of third layer is correct.
@@ -582,16 +559,16 @@ class TestMetricTensor:
 
         layer3_diag = qml.QNode(layer3_diag, dev)
         G3 = layer3_diag(x, y, z, h, g, f) / 4
-        assert np.allclose(G[3:4, 3:4], G3, atol=tol, rtol=0)
-        assert np.allclose(G_alias[3:4, 3:4], G3, atol=tol, rtol=0)
+        assert qml.math.allclose(G[3:4, 3:4], G3, atol=tol, rtol=0)
+        assert qml.math.allclose(G_alias[3:4, 3:4], G3, atol=tol, rtol=0)
 
         # ============================================
         # Finally, double check that the entire metric
         # tensor is as computed.
 
         G_expected = block_diag(G1, G3, G2)
-        assert np.allclose(G, G_expected, atol=tol, rtol=0)
-        assert np.allclose(G_alias, G_expected, atol=tol, rtol=0)
+        assert qml.math.allclose(G, G_expected, atol=tol, rtol=0)
+        assert qml.math.allclose(G_alias, G_expected, atol=tol, rtol=0)
 
     def test_multi_qubit_gates(self):
         """Test that a tape with Ising gates has the correct metric tensor tapes."""
@@ -618,35 +595,435 @@ class TestMetricTensor:
         assert [[type(op) for op in tape.operations] for tape in tapes] == expected_ops
 
 
-@pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop"])
+fixed_pars = np.array([-0.2, 0.2, 0.5, 0.3, 0.7], requires_grad=False)
+
+
+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)
+    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)
+    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(x, wires=None):
+    qml.RX(fixed_pars[0], wires=0)
+    qml.RX(x, wires=0)
+
+
+def fubini_ansatz8(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])
+
+
+fubini_ansatze = [
+    fubini_ansatz0,
+    fubini_ansatz1,
+    fubini_ansatz2,
+    fubini_ansatz3,
+    fubini_ansatz4,
+    fubini_ansatz5,
+    fubini_ansatz6,
+    fubini_ansatz7,
+    fubini_ansatz8,
+]
+
+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,
+        ),
+    ),
+    (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, -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.1111, 0.3333], requires_grad=True),),
+]
+
+
+def autodiff_metric_tensor(ansatz, num_wires):
+    """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 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)
+
+        jac = rjac + 1j * ijac
+        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], [0]))
+        return np.real(
+            np.tensordot(np.conj(jac), jac, axes=([0], [0]))
+            - np.tensordot(np.conj(psidpsi), psidpsi, axes=0)
+        )
+
+    return mt
+
+
+class TestFullMetricTensor:
+
+    num_wires = 3
+
+    @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)
+        dev = qml.device("default.qubit.autograd", wires=self.num_wires + 1)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires[:-1])
+            return qml.expval(qml.PauliZ(0))
+
+        mt = qml.metric_tensor(circuit, approx=None)(*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.skip(reason="JAX does not support the forward pass metric tensor.")
+    @pytest.mark.parametrize("ansatz, params", zip(fubini_ansatze, fubini_params))
+    def test_correct_output_jax(self, ansatz, params):
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
+
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        dev = qml.device("default.qubit.jax", wires=self.num_wires + 1)
+
+        params = tuple(jnp.array(p) for p in params)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires[:-1])
+            return qml.expval(qml.PauliZ(0))
+
+        mt = qml.metric_tensor(circuit, approx=None)(*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", zip(fubini_ansatze, fubini_params))
+    def test_correct_output_torch(self, ansatz, params):
+        torch = pytest.importorskip("torch")
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        dev = qml.device("default.qubit.torch", wires=self.num_wires + 1)
+
+        params = tuple(torch.tensor(p, dtype=torch.float64, requires_grad=True) for p in params)
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires[:-1])
+            return qml.expval(qml.PauliZ(0))
+
+        mt = qml.metric_tensor(circuit, approx=None)(*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", zip(fubini_ansatze, fubini_params))
+    def test_correct_output_tf(self, ansatz, params):
+        tf = pytest.importorskip("tensorflow")
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
+        dev = qml.device("default.qubit.tf", wires=self.num_wires + 1)
+
+        params = tuple(tf.Variable(p, dtype=tf.float64) for p in params)
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(*params):
+            """Circuit with dummy output to create a QNode."""
+            ansatz(*params, dev.wires[:-1])
+            return qml.expval(qml.PauliZ(0))
+
+        with tf.GradientTape() as t:
+            qml.metric_tensor(circuit, approx="block-diag")(*params)
+            mt = qml.metric_tensor(circuit, approx=None)(*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 diffability_ansatz_0(weights, wires=None):
+    qml.RX(weights[0], wires=0)
+    qml.RX(weights[1], wires=0)
+    qml.CNOT(wires=[0, 1])
+    qml.RZ(weights[2], wires=1)
+
+
+expected_diag_jac_0 = lambda weights: np.array(
+    [
+        [0, 0, 0],
+        [0, 0, 0],
+        [
+            np.cos(weights[0] + weights[1]) * np.sin(weights[0] + weights[1]) / 2,
+            np.cos(weights[0] + weights[1]) * np.sin(weights[0] + weights[1]) / 2,
+            0,
+        ],
+    ]
+)
+
+
+def diffability_ansatz_1(weights, wires=None):
+    qml.RX(weights[0], wires=0)
+    qml.RY(weights[1], wires=0)
+    qml.CNOT(wires=[0, 1])
+    qml.RZ(weights[2], wires=1)
+
+
+expected_diag_jac_1 = lambda weights: np.array(
+    [
+        [0, 0, 0],
+        [-np.sin(2 * weights[0]) / 4, 0, 0],
+        [
+            np.cos(weights[0]) * np.cos(weights[1]) ** 2 * np.sin(weights[0]) / 2,
+            np.cos(weights[0]) ** 2 * np.sin(2 * weights[1]) / 4,
+            0,
+        ],
+    ]
+)
+
+
+def diffability_ansatz_2(weights, wires=None):
+    qml.RX(weights[0], wires=0)
+    qml.RY(weights[1], wires=1)
+    qml.CNOT(wires=[0, 1])
+    qml.RZ(weights[2], wires=1)
+
+
+expected_diag_jac_2 = lambda weights: np.array(
+    [
+        [0, 0, 0],
+        [0, 0, 0],
+        [
+            np.cos(weights[1]) ** 2 * np.sin(2 * weights[0]) / 4,
+            np.cos(weights[0]) ** 2 * np.sin(2 * weights[1]) / 4,
+            0,
+        ],
+    ]
+)
+
+weights = np.array([0.432, 0.12, -0.292], requires_grad=True)
+
+
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
+@pytest.mark.parametrize(
+    "ansatz, weights, expected_diag_jac",
+    [
+        (diffability_ansatz_0, (weights,), expected_diag_jac_0),
+        (diffability_ansatz_1, (weights,), expected_diag_jac_1),
+        (diffability_ansatz_2, (weights,), expected_diag_jac_2),
+    ],
+)
 class TestDifferentiability:
     """Test for metric tensor differentiability"""
 
-    def test_autograd(self, diff_method, tol):
-        """Test metric tensor differentiability in the autograd interface"""
-        dev = qml.device("default.qubit", wires=2)
-
-        @qml.qnode(dev, interface="autograd", diff_method=diff_method)
-        def circuit(weights):
-            qml.RX(weights[0], wires=0)
-            qml.RY(weights[1], wires=0)
-            qml.CNOT(wires=[0, 1])
-            qml.PhaseShift(weights[2], wires=1)
+    def get_circuit(self, ansatz):
+        def circuit(*weights):
+            ansatz(*weights)
             return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
 
-        def cost(weights):
-            return qml.metric_tensor(circuit, approx="block-diag")(weights)[2, 2]
+        return circuit
+
+    dev = qml.device("default.qubit", wires=3)
+
+    def test_autograd_diag(self, diff_method, tol, ansatz, weights, expected_diag_jac):
+        """Test metric tensor differentiability in the autograd interface"""
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="autograd", diff_method=diff_method)
+        qnode(*weights)
+
+        def cost_diag(*weights):
+            mt = qml.metric_tensor(qnode, approx="block-diag")(*weights)
+            if isinstance(mt, tuple):
+                diag = qml.math.hstack(
+                    [qml.math.diag(_mt) if len(qml.math.shape(_mt)) == 2 else _mt for _mt in mt]
+                )
+
+            else:
+                diag = qml.math.diag(mt)
+            return diag
+
+        jac = qml.jacobian(cost_diag)(*weights)
+        if isinstance(jac, tuple):
+            assert all(
+                qml.math.allclose(j, e, atol=tol, rtol=0)
+                for j, e in zip(jac, expected_diag_jac(*weights))
+            )
+        else:
+            assert qml.math.allclose(jac, expected_diag_jac(*weights), atol=tol, rtol=0)
+
+    def test_autograd(self, diff_method, tol, ansatz, weights, expected_diag_jac):
+        """Test metric tensor differentiability in the autograd interface"""
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="autograd", diff_method=diff_method)
 
-        weights = np.array([0.432, 0.12, -0.432], requires_grad=True)
-        a, b, c = weights
+        def cost_full(*weights):
+            return np.array(qml.metric_tensor(qnode, approx=None)(*weights))
 
-        grad = qml.grad(cost)(weights)
-        expected = np.array(
-            [np.cos(a) * np.cos(b) ** 2 * np.sin(a) / 2, np.cos(a) ** 2 * np.sin(2 * b) / 4, 0]
+        _cost_full = lambda *weights: np.array(
+            autodiff_metric_tensor(ansatz, num_wires=3)(*weights)
         )
-        assert np.allclose(grad, expected, atol=tol, rtol=0)
+        _c = _cost_full(*weights)
+        c = cost_full(*weights)
+        assert all(
+            qml.math.allclose(_sub_c, sub_c, atol=tol, rtol=0) for _sub_c, sub_c in zip(_c, c)
+        )
+        for argnum in range(len(weights)):
+            expected_full = qml.jacobian(_cost_full, argnum=argnum)(*weights)
+            jac = qml.jacobian(cost_full, argnum=argnum)(*weights)
+            assert qml.math.allclose(expected_full, jac, atol=tol, rtol=0)
 
-    def test_jax(self, diff_method, tol):
+    def test_jax_diag(self, diff_method, tol, ansatz, weights, expected_diag_jac):
         """Test metric tensor differentiability in the JAX interface"""
         if diff_method == "parameter-shift":
             pytest.skip("Does not support parameter-shift")
@@ -654,83 +1031,310 @@ class TestDifferentiability:
         jax = pytest.importorskip("jax")
         from jax import numpy as jnp
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="jax", diff_method=diff_method)
 
-        @qml.qnode(dev, interface="jax", diff_method="backprop")
-        def circuit(weights):
-            qml.RX(weights[0], wires=0)
-            qml.RY(weights[1], wires=0)
-            qml.CNOT(wires=[0, 1])
-            qml.PhaseShift(weights[2], wires=1)
-            return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
+        def cost_diag(*weights):
+            return jnp.diag(qml.metric_tensor(qnode, approx="block-diag")(*weights))
 
-        circuit.interface = "jax"
+        jac = jax.jacobian(cost_diag)(jnp.array(*weights))
+        assert qml.math.allclose(jac, expected_diag_jac(*weights), atol=tol, rtol=0)
 
-        def cost(weights):
-            return qml.metric_tensor(circuit, approx="block-diag")(weights)[2, 2]
+    def test_jax(self, diff_method, tol, ansatz, weights, expected_diag_jac):
+        """Test metric tensor differentiability in the JAX interface"""
+        if diff_method == "parameter-shift":
+            pytest.skip("Does not support parameter-shift")
 
-        weights = jnp.array([0.432, 0.12, -0.432])
-        a, b, c = weights
+        jax = pytest.importorskip("jax")
+        from jax import numpy as jnp
 
-        grad = jax.grad(cost)(weights)
-        expected = np.array(
-            [np.cos(a) * np.cos(b) ** 2 * np.sin(a) / 2, np.cos(a) ** 2 * np.sin(2 * b) / 4, 0]
-        )
-        assert np.allclose(grad, expected, atol=tol, rtol=0)
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="jax", diff_method=diff_method)
+
+        def cost_full(*weights):
+            return qml.metric_tensor(qnode, approx=None)(*weights)
 
-    def test_tf(self, diff_method, tol):
+        _cost_full = lambda *weights: autodiff_metric_tensor(ansatz, num_wires=3)(*weights)
+        assert qml.math.allclose(_cost_full(*weights), cost_full(*weights), atol=tol, rtol=0)
+        jac = jax.jacobian(cost_full)(*weights)
+        expected_full = qml.jacobian(_cost_full)(*weights)
+        assert qml.math.allclose(expected_full, jac, atol=tol, rtol=0)
+
+    def test_tf_diag(self, diff_method, tol, ansatz, weights, expected_diag_jac):
         """Test metric tensor differentiability in the TF interface"""
         tf = pytest.importorskip("tensorflow", minversion="2.0")
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="tf", diff_method=diff_method)
 
-        dev = qml.device("default.qubit", wires=2)
-
-        @qml.qnode(dev, interface="tf", diff_method=diff_method)
-        def circuit(weights):
-            qml.RX(weights[0], wires=0)
-            qml.RY(weights[1], wires=0)
-            qml.CNOT(wires=[0, 1])
-            qml.PhaseShift(weights[2], wires=1)
-            return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
+        weights_t = tuple(tf.Variable(w) for w in weights)
+        with tf.GradientTape() as tape:
+            loss_diag = tf.linalg.diag_part(
+                qml.metric_tensor(qnode, approx="block-diag")(*weights_t)
+            )
+        jac = tape.jacobian(loss_diag, weights_t)
+        assert qml.math.allclose(jac, expected_diag_jac(*weights), atol=tol, rtol=0)
 
-        weights = np.array([0.432, 0.12, -0.432])
-        weights_t = tf.Variable(weights)
-        a, b, c = weights
+    def test_tf(self, diff_method, tol, ansatz, weights, expected_diag_jac):
+        """Test metric tensor differentiability in the TF interface"""
+        tf = pytest.importorskip("tensorflow", minversion="2.0")
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="tf", diff_method=diff_method)
 
+        weights_t = tuple(tf.Variable(w) for w in weights)
         with tf.GradientTape() as tape:
-            loss = qml.metric_tensor(circuit, approx="block-diag")(weights_t)[2, 2]
+            loss_full = qml.metric_tensor(qnode, approx=None)(*weights_t)
+        jac = tape.jacobian(loss_full, weights_t)
+        _cost_full = lambda *weights: autodiff_metric_tensor(ansatz, num_wires=3)(*weights)
+        assert qml.math.allclose(_cost_full(*weights), loss_full, atol=tol, rtol=0)
+        expected_full = qml.jacobian(_cost_full)(*weights)
+        assert qml.math.allclose(expected_full, jac, atol=tol, rtol=0)
+
+    def test_torch_diag(self, diff_method, tol, ansatz, weights, expected_diag_jac):
+        """Test metric tensor differentiability in the torch interface"""
+        torch = pytest.importorskip("torch")
 
-        grad = tape.gradient(loss, weights_t)
-        expected = np.array(
-            [np.cos(a) * np.cos(b) ** 2 * np.sin(a) / 2, np.cos(a) ** 2 * np.sin(2 * b) / 4, 0]
-        )
-        assert np.allclose(grad, expected, atol=tol, rtol=0)
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="torch", diff_method=diff_method)
+
+        weights_t = tuple(torch.tensor(w, requires_grad=True) for w in weights)
+
+        def cost_diag(*weights):
+            mt = qml.metric_tensor(qnode, approx="block-diag")(*weights)
+            if isinstance(mt, tuple):
+                diag = qml.math.hstack(
+                    [qml.math.diag(_mt) if len(qml.math.shape(_mt)) == 2 else _mt for _mt in mt]
+                )
+
+            else:
+                diag = qml.math.diag(mt)
+            return diag
+
+        jac = torch.autograd.functional.jacobian(cost_diag, weights_t)
 
-    def test_torch(self, diff_method, tol):
+        if isinstance(jac, tuple) and len(jac) != 1:
+            assert all(
+                qml.math.allclose(j.detach().numpy(), e, atol=tol, rtol=0)
+                for j, e in zip(jac, expected_diag_jac(*weights))
+            )
+        else:
+            if isinstance(jac, tuple) and len(jac) == 1:
+                jac = jac[0]
+            assert qml.math.allclose(
+                jac.detach().numpy(), expected_diag_jac(*weights), atol=tol, rtol=0
+            )
+
+    def test_torch(self, diff_method, tol, ansatz, weights, expected_diag_jac):
         """Test metric tensor differentiability in the torch interface"""
         torch = pytest.importorskip("torch")
 
-        dev = qml.device("default.qubit", wires=2)
+        circuit = self.get_circuit(ansatz)
+        qnode = qml.QNode(circuit, self.dev, interface="torch", diff_method=diff_method)
+        weights_t = tuple(torch.tensor(w, requires_grad=True) for w in weights)
+        qnode(*weights_t)
+        cost_full = qml.metric_tensor(qnode, approx=None)
+        _cost_full = autodiff_metric_tensor(ansatz, num_wires=3)
+        jac = torch.autograd.functional.jacobian(cost_full, weights_t)
+        expected_full = qml.jacobian(_cost_full)(*weights)
+        assert qml.math.allclose(
+            _cost_full(*weights), cost_full(*weights_t).detach().numpy(), atol=tol, rtol=0
+        )
+        if isinstance(jac, tuple) and len(jac) != 1:
+            assert all(
+                qml.math.allclose(j.detach().numpy(), e, atol=tol, rtol=0)
+                for j, e in zip(jac, expected_full)
+            )
+        else:
+            if isinstance(jac, tuple) and len(jac) == 1:
+                jac = jac[0]
+            assert qml.math.allclose(jac.detach().numpy(), expected_full, atol=tol, rtol=0)
 
-        @qml.qnode(dev, interface="torch", diff_method=diff_method)
-        def circuit(weights):
-            qml.RX(weights[0], wires=0)
-            qml.RY(weights[1], wires=0)
-            qml.CNOT(wires=[0, 1])
-            qml.PhaseShift(weights[2], wires=1)
-            return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
 
-        weights = np.array([0.432, 0.12, -0.432])
-        a, b, c = weights
+@pytest.mark.parametrize("approx", [True, False, "Invalid", 2])
+def test_invalid_value_for_approx(approx):
+    """Test exception is raised if ``approx`` is invalid."""
+    with qml.tape.QuantumTape() as tape:
+        qml.RX(np.array(0.5, requires_grad=True), wires=0)
+        qml.expval(qml.PauliX(0))
 
-        weights_t = torch.tensor(weights, requires_grad=True)
-        loss = qml.metric_tensor(circuit, approx="block-diag")(weights_t)[2, 2]
-        loss.backward()
+    with pytest.raises(ValueError, match="keyword argument approx"):
+        qml.metric_tensor(tape, approx=approx)
 
-        grad = weights_t.grad
-        expected = np.array(
-            [np.cos(a) * np.cos(b) ** 2 * np.sin(a) / 2, np.cos(a) ** 2 * np.sin(2 * b) / 4, 0]
-        )
-        assert np.allclose(grad, expected, atol=tol, rtol=0)
+
+def test_generator_no_expval(monkeypatch):
+    """Test exception is raised if subcircuit contains an
+    operation with generator object that is not an observable"""
+    with monkeypatch.context() as m:
+        m.setattr("pennylane.RX.generator", [qml.RX, 1])
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(np.array(0.5, requires_grad=True), wires=0)
+            qml.expval(qml.PauliX(0))
+
+        with pytest.raises(qml.QuantumFunctionError, match="no corresponding observable"):
+            qml.metric_tensor(tape, approx="block-diag")
+
+
+def test_error_missing_aux_wire():
+    """Tests that a special error is raised if the requested (or default, if not given)
+    auxiliary wire for the Hadamard test is missing."""
+    dev = qml.device("default.qubit", wires=qml.wires.Wires(["wire1", "wire2"]))
+
+    @qml.qnode(dev)
+    def circuit(x, z):
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire2")
+        qml.CNOT(wires=["wire1", "wire2"])
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire2")
+        return qml.expval(qml.PauliZ("wire2"))
+
+    x = np.array(0.5, requires_grad=True)
+    z = np.array(0.1, requires_grad=True)
+
+    with pytest.warns(UserWarning, match="The device does not have a wire that is not used"):
+        qml.metric_tensor(circuit, approx=None)(x, z)
+    with pytest.warns(UserWarning, match="The device does not have a wire that is not used"):
+        qml.metric_tensor(circuit, approx=None, aux_wire="wire3")(x, z)
+
+
+def test_error_aux_wire_replaced():
+    """Tests that even if an aux_wire is provided, it is superseded by a device
+    wire if it does not exist itself on the device, so that the metric_tensor is
+    successfully computed."""
+    dev = qml.device("default.qubit", wires=qml.wires.Wires(["wire1", "wire2", "hidden_wire"]))
+
+    @qml.qnode(dev)
+    def circuit(x, z):
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire2")
+        qml.CNOT(wires=["wire1", "wire2"])
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire2")
+        return qml.expval(qml.PauliZ("wire2"))
+
+    x = np.array(0.5, requires_grad=True)
+    z = np.array(0.1, requires_grad=True)
+
+    qml.metric_tensor(circuit, approx=None, aux_wire="wire3")(x, z)
+
+
+@pytest.mark.parametrize("allow_nonunitary", [True, False])
+def test_error_generator_not_registered(allow_nonunitary, monkeypatch):
+    """Tests that an error is raised if an operation doe not have a
+    controlled-generator operation registered."""
+    dev = qml.device("default.qubit", wires=qml.wires.Wires(["wire1", "wire2", "wire3"]))
+
+    x = np.array(0.5, requires_grad=True)
+    z = np.array(0.1, requires_grad=True)
+
+    @qml.qnode(dev)
+    def circuit(x, z):
+        qml.CRX(x, wires=["wire1", "wire2"])
+        qml.RZ(z, wires="wire2")
+        return qml.expval(qml.PauliZ("wire2"))
+
+    with monkeypatch.context() as m:
+        exp_fn = lambda tape, *args, **kwargs: tape
+        m.setattr("pennylane.transforms.metric_tensor.expand_fn", exp_fn)
+
+        if allow_nonunitary:
+            qml.metric_tensor(circuit, approx=None, allow_nonunitary=allow_nonunitary)(x, z)
+        else:
+            with pytest.raises(ValueError, match="Generator for operation"):
+                qml.metric_tensor(circuit, approx=None, allow_nonunitary=allow_nonunitary)(x, z)
+
+    @qml.qnode(dev)
+    def circuit(x, z):
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire1")
+        return qml.expval(qml.PauliZ("wire2"))
+
+    with monkeypatch.context() as m:
+        exp_fn = lambda tape, *args, **kwargs: tape
+        m.setattr("pennylane.transforms.metric_tensor.expand_fn", exp_fn)
+        m.setattr("pennylane.RX.generator", [qml.Hadamard, 1.0])
+
+        if allow_nonunitary:
+            qml.metric_tensor(circuit, approx=None, allow_nonunitary=allow_nonunitary)(x, z)
+        else:
+            with pytest.raises(ValueError, match="Generator for operation"):
+                qml.metric_tensor(circuit, approx=None, allow_nonunitary=allow_nonunitary)(x, z)
+
+
+def test_no_error_missing_aux_wire_not_used():
+    """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."""
+    dev = qml.device("default.qubit", wires=qml.wires.Wires(["wire1", "wire2"]))
+
+    @qml.qnode(dev)
+    def circuit_single_block(x, z):
+        """This circuit has a metric tensor that consists
+        of a single block in the block diagonal "approximation"."""
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire2")
+        qml.CNOT(wires=["wire1", "wire2"])
+        return qml.expval(qml.PauliZ("wire2"))
+
+    @qml.qnode(dev)
+    def circuit_multi_block(x, z):
+        """This circuit has a metric tensor that consists
+        of multiple blocks and thus is approximated when only
+        computing the block diagonal."""
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire2")
+        qml.CNOT(wires=["wire1", "wire2"])
+        qml.RX(x, wires="wire1")
+        qml.RZ(z, wires="wire2")
+        return qml.expval(qml.PauliZ("wire2"))
+
+    x = np.array(0.5, requires_grad=True)
+    z = np.array(0.1, requires_grad=True)
+
+    qml.metric_tensor(circuit_single_block, approx=None)(x, z)
+    qml.metric_tensor(circuit_single_block, approx=None, aux_wire="aux_wire")(x, z)
+    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)
+
+
+def aux_wire_ansatz_0(x, y):
+    qml.RX(x, wires=0)
+    qml.RY(x, wires=2)
+
+
+def aux_wire_ansatz_1(x, y):
+    qml.RX(x, wires=0)
+    qml.RY(x, wires=1)
+
+
+@pytest.mark.parametrize("aux_wire", [None, "aux", 3])
+@pytest.mark.parametrize("ansatz", [aux_wire_ansatz_0, aux_wire_ansatz_1])
+def test_get_aux_wire(aux_wire, ansatz):
+    """Test ``_get_aux_wire`` without device_wires."""
+    x, y = np.array([0.2, 0.1], requires_grad=True)
+    with qml.tape.JacobianTape() as tape:
+        ansatz(x, y)
+    out = _get_aux_wire(aux_wire, tape, None)
+
+    if aux_wire is not None:
+        assert out == aux_wire
+    else:
+        assert out == (1 if 1 not in tape.wires else 2)
+
+
+def test_get_aux_wire_with_device_wires():
+    """Test ``_get_aux_wire`` with device_wires."""
+    x, y = np.array([0.2, 0.1], requires_grad=True)
+    with qml.tape.JacobianTape() as tape:
+        qml.RX(x, wires=0)
+        qml.RX(x, wires="one")
+
+    device_wires = qml.wires.Wires([0, "aux", "one"])
+
+    assert _get_aux_wire(0, tape, device_wires) == 0
+    assert _get_aux_wire("one", tape, device_wires) == "one"
+    assert _get_aux_wire("two", tape, device_wires) == "aux"
+    assert _get_aux_wire(None, tape, device_wires) == "aux"
 
 
 class TestDeprecatedQNodeMethod:
@@ -765,7 +1369,7 @@ class TestDeprecatedQNodeMethod:
             )
             / 4
         )
-        assert np.allclose(g, np.diag(expected), atol=tol, rtol=0)
+        assert qml.math.allclose(g, np.diag(expected), atol=tol, rtol=0)
 
     def test_tapes_returned(self, tol):
         """Test that a warning is emitted"""