Repository: PennyLaneAI/pennylane (python)

Issue:
QNGOptimizer returns TypeError when step method called 
* *Expected behavior:* Returns a `list` of updated parameters

* *Actual behavior:* Returns a `TypeError`

* *System information:* 
Platform info:           Linux-5.8.0-45-generic-x86_64-with-glibc2.29
Python version:          3.8.5
Numpy version:           1.19.5
Scipy version:           1.4.1

#### Source code and tracebacks
```python
import networkx as nx
import pennylane as qml

edges = [(0, 1), (1, 2), (2, 0), (2, 3)]
graph = nx.Graph(edges)

H_c, H_m = qml.qaoa.min_vertex_cover(graph, constrained=False)

def qaoa_layer(gamma, alpha):
    qml.qaoa.cost_layer(gamma, H_c)
    qml.qaoa.mixer_layer(alpha, H_m)
    
depth = 5

def circuit(params, **kwargs):
    
    qml.PauliX(wires=0)  
    qml.PauliX(wires=3)
    
    qml.layer(qaoa_layer, depth, params[0], params[1])

dev = qml.device("default.qubit", wires=4)
cost_function = qml.ExpvalCost(circuit, H_c, dev, optimize=True)
params = [[0.5]*depth,[0.5]*depth]

optimizer = qml.QNGOptimizer(stepsize=0.1)


for _ in range(100):
    params = optimizer.step(cost_function, params)

```




Pull request description:
Add support for batch execution to `qml.metric_tensor`
**Context:** With the low-level differentiable batch-execution pipeline now available in PennyLane, and the availability of `@qml.batch_transform`, we can now start porting our existing batch-tape transforms to take advantage of this framework.

**Description of the Change:**

- When using `qml.math.get_trainable_indices()`, only NumPy arrays with `requires_grad=True` are taken as trainable when using Autograd. This is a requirement for batch transformations like the metric tensor to work correctly.

- When using the QNode in backpropagation, trainable parameter indices are computed and stored. Previously, a backpropagation QNode would simply mark _all_ parameters as trainable. This extra information is needed for the metric tensor.

- The metric tensor has been converted into a batch transformation, that accepts both tapes _and_ QNodes as input. This makes the old `metric_tensor_tape` function irrelevant, and it has been removed. The tests have been likewise updated.

- Previously, `qml.metric_tensor(qnode)(*args, **kwargs)` would only return the metric tensor with respect to gate arguments, and ignore any classical processing inside the QNode, even very trivial classical processing such as parameter permutation. This would lead to many reported user bugs, such as #1154. In this new framework, the metric tensor now takes into account classical processing, and returns the metric tensor with respect to _QNode arguments_, not simply _gate_ arguments.

  To revert to the previous behaviour of returning the metric tensor with respect to gate arguments, `qml.metric_tensor(qnode, hybrid=False)` can be passed.

**Benefits:**

- The `qml.metric_tensor()` function now makes use of batch execution for submission of circuits required for metric tensor computation.

- The `qml.metric_tensor()` function now takes into account classical computation inside the QNode, for example, if gate arguments are repeated or permuted.

- QNodes in backpropagation mode correctly track trainable parameters, fixing an issue that was long part of our test suite.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** Closes #1154 


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 12d08892c..5ba511c52 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -40,11 +40,60 @@
 
 <h3>Improvements</h3>
 
+* The `qml.metric_tensor` transform has been improved with regards to
+  both function and performance.
+  [(#1638)](https://github.com/PennyLaneAI/pennylane/pull/1638)
+
+  - If the underlying device supports batch execution of circuits, the quantum circuits required to
+    compute the metric tensor elements will be automatically submitted as a batched job. This can
+    lead to significant performance improvements for devices with a non-trivial job submission
+    overhead.
+
+  - Previously, the transform would only return the metric tensor with respect to gate arguments,
+    and ignore any classical processing inside the QNode, even very trivial classical processing
+    such as parameter permutation. The metric tensor now takes into account classical processing,
+    and returns the metric tensor with respect to QNode arguments, not simply gate arguments:
+
+    ```pycon
+    >>> @qml.qnode(dev)
+    ... def circuit(x):
+    ...     qml.Hadamard(wires=1)
+    ...     qml.RX(x[0], wires=0)
+    ...     qml.CNOT(wires=[0, 1])
+    ...     qml.RY(x[1] ** 2, wires=1)
+    ...     qml.RY(x[1], wires=0)
+    ...     return qml.expval(qml.PauliZ(0))
+    >>> x = np.array([0.1, 0.2], requires_grad=True)
+    >>> qml.metric_tensor(circuit)(x)
+    array([[0.25      , 0.        ],
+           [0.        , 0.28750832]])
+    ```
+
+    To revert to the previous behaviour of returning the metric tensor with respect to gate
+    arguments, `qml.metric_tensor(qnode, hybrid=False)` can be passed.
+
+    ```pycon
+    >>> qml.metric_tensor(circuit, hybrid=False)(x)
+    array([[0.25      , 0.        , 0.        ],
+           [0.        , 0.25      , 0.        ],
+           [0.        , 0.        , 0.24750832]])
+    ```
+
 * ``qml.circuit_drawer.CircuitDrawer`` can accept a string for the ``charset`` keyword, instead of a ``CharSet`` object.
   [(#1640)](https://github.com/PennyLaneAI/pennylane/pull/1640)
 
 <h3>Breaking changes</h3>
 
+- The `QNode.metric_tensor` method has been deprecated, and will be removed in an upcoming release.
+  Please use the `qml.metric_tensor` transform instead.
+  [(#1638)](https://github.com/PennyLaneAI/pennylane/pull/1638)
+
+- The utility function `qml.math.requires_grad` now returns `True` when using Autograd
+  if and only if the `requires_grad=True` attribute is set on the NumPy array. Previously,
+  this function would return `True` for *all* NumPy arrays and Python floats, unless
+  `requires_grad=False` was explicitly set.
+  [(#1638)](https://github.com/PennyLaneAI/pennylane/pull/1638)
+
 <h3>Bug fixes</h3>
 
 * The device suite tests can now execute successfully if no shots configuration variable is given.
@@ -56,7 +105,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Christina Lee, David Wierichs
+Josh Izaac, Christina Lee, David Wierichs.
 
 # Release 0.18.0 (Current release)
 
diff --git a/pennylane/math/utils.py b/pennylane/math/utils.py
index f66c375c7..585c930d3 100644
--- a/pennylane/math/utils.py
+++ b/pennylane/math/utils.py
@@ -268,9 +268,7 @@ def requires_grad(tensor, interface=None):
         if isinstance(tensor, ArrayBox):
             return True
 
-        # Currently, in the Autograd interface, we assume
-        # that all objects are differentiable by default.
-        return getattr(tensor, "requires_grad", True)
+        return getattr(tensor, "requires_grad", False)
 
     if interface == "torch":
         return getattr(tensor, "requires_grad", False)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index d718042e1..f088ebeb2 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -644,6 +644,10 @@ class QNode:
         # provide the jacobian options
         self.qtape.jacobian_options = self.diff_options
 
+        if self.diff_options["method"] == "backprop":
+            params = self.qtape.get_parameters(trainable_only=False)
+            self.qtape.trainable_params = qml.math.get_trainable_indices(params)
+
     def __call__(self, *args, **kwargs):
 
         # If shots specified in call but not in qfunc signature,
@@ -725,10 +729,19 @@ class QNode:
         Returns:
             array[float]: metric tensor
         """
-        return qml.metric_tensor(self, diag_approx=diag_approx, only_construct=only_construct)(
-            *args, **kwargs
+        warnings.warn(
+            "The QNode.metric_tensor method has been deprecated. "
+            "Please use the qml.metric_tensor transform instead.",
+            UserWarning,
         )
 
+        if only_construct:
+            self.construct(args, kwargs)
+            tape = qml.metric_tensor.expand_fn(self.qtape)
+            return qml.metric_tensor.construct(tape, diag_approx=diag_approx)
+
+        return qml.metric_tensor(self, diag_approx=diag_approx)(*args, **kwargs)
+
     def draw(
         self, charset="unicode", wire_order=None, show_all_wires=False
     ):  # pylint: disable=unused-argument
diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py
index 986195e95..f267b1396 100644
--- a/pennylane/transforms/__init__.py
+++ b/pennylane/transforms/__init__.py
@@ -82,7 +82,6 @@ more tapes as well as a classical processing function.
     :toctree: api
 
     ~transforms.measurement_grouping
-    ~transforms.metric_tensor_tape
     ~transforms.hamiltonian_expand
 
 Decorators and utility functions
@@ -111,7 +110,7 @@ from .draw import draw
 from .hamiltonian_expand import hamiltonian_expand
 from .invisible import invisible
 from .measurement_grouping import measurement_grouping
-from .metric_tensor import metric_tensor, metric_tensor_tape
+from .metric_tensor import metric_tensor
 from .optimization import (
     cancel_inverses,
     commute_controlled,
diff --git a/pennylane/transforms/classical_jacobian.py b/pennylane/transforms/classical_jacobian.py
index e377fc9e3..0e853d547 100644
--- a/pennylane/transforms/classical_jacobian.py
+++ b/pennylane/transforms/classical_jacobian.py
@@ -135,8 +135,9 @@ def classical_jacobian(qnode, argnum=None):
 
     def classical_preprocessing(*args, **kwargs):
         """Returns the trainable gate parameters for a given QNode input."""
+        trainable_only = kwargs.pop("_trainable_only", True)
         qnode.construct(args, kwargs)
-        return qml.math.stack(qnode.qtape.get_parameters())
+        return qml.math.stack(qnode.qtape.get_parameters(trainable_only=trainable_only))
 
     if qnode.interface == "autograd":
 
@@ -176,6 +177,7 @@ def classical_jacobian(qnode, argnum=None):
         argnum = 0 if argnum is None else argnum
 
         def _jacobian(*args, **kwargs):
+            kwargs["_trainable_only"] = False
             return jax.jacobian(classical_preprocessing, argnums=argnum)(*args, **kwargs)
 
         return _jacobian
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index d84dd7016..0f4ca9b21 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -14,105 +14,141 @@
 """
 Contains the metric tensor transform
 """
+import functools
 import warnings
 
 import numpy as np
 import pennylane as qml
 
 
+from .batch_transform import batch_transform
+
+
+SUPPORTED_OPS = ["RX", "RY", "RZ", "PhaseShift"]
+
+
 def _stopping_critera(obj):
-    if getattr(obj, "num_params", 0) == 0:
-        return True
+    return getattr(obj, "num_params", 0) == 0 or obj.name in SUPPORTED_OPS
+
 
-    if obj.name in ["RX", "RY", "RZ", "PhaseShift"]:
-        return True
+def expand_fn(tape):
+    """Expands the tape to contain only operations
+    supported by the ``metric_tensor`` transform (specified
+    by ``SUPPORTED_OPS``).
+    """
+    new_tape = tape.expand(depth=2, stop_at=_stopping_critera)
+    params = new_tape.get_parameters(trainable_only=False)
+    new_tape.trainable_params = qml.math.get_trainable_indices(params)
+    return new_tape
 
-    return False
 
+@functools.partial(batch_transform, expand_fn=expand_fn)
+def metric_tensor(tape, diag_approx=False):
+    """Returns a function that computes the block-diagonal approximation of the metric tensor
+    of a given QNode or quantum tape.
 
-def metric_tensor_tape(tape, diag_approx=False, wrt=None):
-    """Returns a list of tapes, and a classical processing function, for computing the block
-    diagonal metric tensor approximation of an input tape on hardware.
+    .. note::
+
+        Currently, only the :class:`~.RX`, :class:`~.RY`, :class:`~.RZ`, and
+        :class:`~.PhaseShift` parametrized gates are supported.
+        All other parametrized gates will be decomposed if possible.
 
     Args:
-        tape (.QuantumTape): the tape to compute the metric tensor of
-        diag_approx (bool): If ``True`` the diagonal approximation to the metric
-            tensor is computed. If ``False``, a block diagonal approximation
-            to the metric tensor is computed.
-        wrt (Sequence[int]): Indices of the tape parameters with which to
-            compute the metric tensor. Parameter indices not included are
-            treated as *fixed* parameters. Defaults to the tape's trainable
-            parameters.
+        qnode (.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.
+        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.
+
+            - 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.
 
     Returns:
-        tuple[list[.QuantumTape], func]: Returns a tuple containing a list of
-        quantum tapes to be evaluated, and a function to be applied to these
-        tape results to compute the metric tensor.
+        func: Function which accepts the same arguments as the QNode. When called, this
+        function will return the metric tensor.
 
     **Example**
 
-    Given the following quantum tape,
+    Consider the following QNode:
 
     .. code-block:: python
 
-        with qml.tape.QuantumTape() as tape:
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(weights):
             # layer 1
-            qml.RX(0.1, wires=0)
-            qml.RX(0.2, wires=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(0.4, wires=0)
-            qml.RZ(0.5, wires=2)
+            qml.RZ(weights[2], wires=0)
+            qml.RZ(weights[3], wires=2)
 
             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.expval(qml.PauliZ(0) @ qml.PauliZ(1))
-            qml.expval(qml.PauliY(2))
-
-    We can use the ``metric_tensor`` transform to generate a new tapes and a classical
-    processing function for computing the metric tensor.
-
-    >>> mt_tapes, fn = qml.transforms.metric_tensor_tape(tape)
-    >>> print(mt_tapes)
-    [<QuantumTape: wires=[0, 1, 2], params=0>, <QuantumTape: wires=[0, 1, 2], params=2>]
-    >>> print(mt_tapes[0].draw())
-     0: ──H──╭┤ Probs
-     1: ──H──├┤ Probs
-     2: ─────╰┤ Probs
-    >>> print(mt_tapes[1].draw())
-     0: ──RX(0.1)──╭C──────╭┤ Probs
-     1: ──RX(0.2)──╰X──╭C──├┤ Probs
-     2: ───────────────╰X──╰┤ Probs
-
-    We can evaluate these tapes on a device:
-
-    >>> dev = qml.device("default.qubit", wires=3)
-    >>> res = dev.batch_execute(mt_tapes)
-    >>> print(res)
-    [array([[0.25, 0.  , 0.25, 0.  , 0.25, 0.  , 0.25, 0.  ]]),
-     array([[9.87560268e-01, 0.00000000e+00, 0.00000000e+00, 9.94181506e-03,
-             2.48960206e-05, 0.00000000e+00, 0.00000000e+00, 2.47302134e-03]])]
-
-    Applying the processing function results in the metric tensor:
-
-    >>> fn(res)
-    array([[0.25      , 0.        , 0.        , 0.        ],
-           [0.        , 0.25      , 0.        , 0.        ],
-           [0.        , 0.        , 0.00249168, 0.00244201],
-           [0.        , 0.        , 0.00244201, 0.01226071]])
-    """
+    We can use the ``metric_tensor`` transform to generate a new function that returns the
+    metric tensor of this QNode:
 
-    # For parametrized operations, only the RX, RY, RZ, and PhaseShift gates are supported.
-    # Expand out all other gates.
-    tape = tape.expand(depth=2, stop_at=_stopping_critera)
+    >>> met_fn = qml.metric_tensor(circuit)
+    >>> weights = np.array([0.1, 0.2, 0.4, 0.5], requires_grad=True)
+    >>> met_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)
+
+    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``:
 
-    if wrt is not None:
-        tape.trainable_params = set(wrt)
+    >>> grad_fn = qml.grad(lambda x: met_fn(x)[3, 2])
+    >>> grad_fn(weights)
+    array([[ 0.04867729, -0.00049502,  0.        ],
+           [ 0.        ,  0.        ,  0.        ]])
 
+    .. UsageDetails::
+
+        This transform can also be applied to low-level
+        :class:`~.QuantumTape` objects. This will result in no implicit quantum
+        device evaluation. Instead, the processed tapes, and post-processing
+        function, which together define the metric tensor are directly returned:
+
+        >>> params = np.array([1.7, 1.0, 0.5], requires_grad=True)
+        >>> with qml.tape.QuantumTape() as tape:
+        ...     qml.RX(params[0], wires=0)
+        ...     qml.RY(params[1], wires=0)
+        ...     qml.CNOT(wires=[0, 1])
+        ...     qml.PhaseShift(params[2], wires=1)
+        ...     qml.expval(qml.PauliX(0))
+        >>> tapes, fn = qml.metric_tensor(tape)
+        >>> tapes
+        [<QuantumTape: wires=[0, 1], params=0>,
+         <QuantumTape: wires=[0, 1], params=1>,
+         <QuantumTape: wires=[0, 1], params=3>]
+
+        This can be useful if the underlying circuits representing the metric tensor
+        computation need to be analyzed.
+
+        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]])
+    """
     # get the circuit graph
     graph = tape.graph
 
@@ -152,7 +188,7 @@ def metric_tensor_tape(tape, diag_approx=False, wrt=None):
         # to measure in the basis of the parametrized layer generators.
         with tape.__class__() as layer_tape:
             for op in queue:
-                op.queue()
+                qml.apply(op)
 
             for o in obs_list[-1]:
                 o.diagonalizing_gates()
@@ -177,71 +213,14 @@ def metric_tensor_tape(tape, diag_approx=False, wrt=None):
     return metric_tensor_tapes, processing_fn
 
 
-def metric_tensor(qnode, diag_approx=False, only_construct=False):
-    """Returns a function that returns the value of the metric tensor
-    of a given QNode.
-
-    .. note::
+@metric_tensor.custom_qnode_wrapper
+def qnode_execution_wrapper(self, qnode, targs, tkwargs):
+    """Here, we overwrite the QNode execution wrapper in order
+    to take into account that classical processing may be present
+    inside the QNode."""
+    hybrid = tkwargs.pop("hybrid", True)
 
-        Currently, only the :class:`~.RX`, :class:`~.RY`, :class:`~.RZ`, and
-        :class:`~.PhaseShift` parametrized gates are supported.
-        All other parametrized gates will be decomposed if possible.
-
-    Args:
-        qnode (.QNode or .ExpvalCost): QNode(s) to compute the metric tensor of
-        diag_approx (bool): iff True, use the diagonal approximation
-        only_construct (bool): Iff True, construct the circuits used for computing
-            the metric tensor but do not execute them, and return the tapes.
-
-    Returns:
-        func: Function which accepts the same arguments as the QNode. When called, this
-        function will return the metric tensor.
-
-    **Example**
-
-    Consider the following QNode:
-
-    .. code-block:: python
-
-        dev = qml.device("default.qubit", wires=3)
-
-        @qml.qnode(dev, interface="autograd")
-        def circuit(weights):
-            # layer 1
-            qml.RX(weights[0, 0], wires=0)
-            qml.RX(weights[0, 1], wires=1)
-
-            qml.CNOT(wires=[0, 1])
-            qml.CNOT(wires=[1, 2])
-
-            # layer 2
-            qml.RZ(weights[1, 0], wires=0)
-            qml.RZ(weights[1, 1], wires=2)
-
-            qml.CNOT(wires=[0, 1])
-            qml.CNOT(wires=[1, 2])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)), qml.expval(qml.PauliY(2))
-
-    We can use the ``metric_tensor`` function to generate a new function, that returns the
-    metric tensor of this QNode:
-
-    >>> met_fn = qml.metric_tensor(circuit)
-    >>> weights = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], requires_grad=True)
-    >>> met_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)
-
-    The returned metric tensor is also fully differentiable, in all interfaces.
-    For example, differentiating the ``(3, 2)`` element:
-
-    >>> grad_fn = qml.grad(lambda x: met_fn(x)[3, 2])
-    >>> grad_fn(weights)
-    array([[ 0.04867729, -0.00049502,  0.        ],
-           [ 0.        ,  0.        ,  0.        ]])
-    """
-    if qnode.__class__.__name__ == "ExpvalCost":
+    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 "
@@ -250,30 +229,43 @@ def metric_tensor(qnode, diag_approx=False, only_construct=False):
 
         qnode = qnode.qnodes.qnodes[0]
 
-    def _metric_tensor_fn(*args, **kwargs):
-        jac = qml.math.stack(qml.transforms.classical_jacobian(qnode)(*args, **kwargs))
-        jac = qml.math.reshape(jac, [qnode.qtape.num_params, -1])
+    mt_fn = self.default_qnode_wrapper(qnode, targs, tkwargs)
+    cjac_fn = qml.transforms.classical_jacobian(qnode)
+
+    def wrapper(*args, **kwargs):
+        mt = mt_fn(*args, **kwargs)
+
+        if not hybrid:
+            return mt
 
-        wrt, perm = np.nonzero(qml.math.toarray(jac))
-        perm = np.argsort(np.argsort(perm))
+        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 mt @ cjac.
+                metric_tensors = []
 
-        qnode.construct(args, kwargs)
+                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]])
+                        metric_tensors.append(_mt)
 
-        metric_tensor_tapes, processing_fn = metric_tensor_tape(
-            qnode.qtape,
-            diag_approx=diag_approx,
-            wrt=wrt.tolist() if qnode.diff_options["method"] == "backprop" else None,
-        )
+                return tuple(metric_tensors)
 
-        if only_construct:
-            return metric_tensor_tapes
+        is_square = cjac.shape == (1,) or (cjac.ndim == 2 and cjac.shape[0] == cjac.shape[1])
 
-        res = [t.execute(device=qnode.device) for t in metric_tensor_tapes]
-        mt = processing_fn(res)
+        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
 
-        # permute rows ad columns
-        mt = qml.math.gather(mt, perm)
-        mt = qml.math.gather(qml.math.T(mt), perm)
+        # 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
 
-    return _metric_tensor_fn
+    return wrapper


Test patch:
diff --git a/tests/gradients/test_gradient_transform.py b/tests/gradients/test_gradient_transform.py
index 260446f4f..5b84db167 100644
--- a/tests/gradients/test_gradient_transform.py
+++ b/tests/gradients/test_gradient_transform.py
@@ -270,7 +270,7 @@ class TestGradientTransformIntegration:
 
         res = qml.gradients.param_shift(circuit)(d, w)
         classical_jac = spy.spy_return(d, w)
-        assert np.allclose(classical_jac, np.array([[0, 2 * w[0], 0], [0, 0, 1]]).T)
+        assert np.allclose(classical_jac, np.array([[2 * w[0], 0], [0, 1]]).T)
 
         expected = np.array([-2 * x * np.cos(np.cos(d)) * np.sin(x ** 2), 0])
         assert np.allclose(res, expected, atol=tol, rtol=0)
diff --git a/tests/interfaces/test_qnode_autograd.py b/tests/interfaces/test_qnode_autograd.py
index df8f94e51..c4d81d054 100644
--- a/tests/interfaces/test_qnode_autograd.py
+++ b/tests/interfaces/test_qnode_autograd.py
@@ -423,22 +423,10 @@ class TestQNode:
             return qml.expval(qml.PauliX(0))
 
         res = circuit(a, p)
-
-        if diff_method == "finite-diff":
-            assert circuit.qtape.trainable_params == {1, 2, 3, 4}
-        elif diff_method == "backprop":
-            # For a backprop device, no interface wrapping is performed, and JacobianTape.jacobian()
-            # is never called. As a result, JacobianTape.trainable_params is never set --- the ML
-            # framework uses its own backprop logic and its own bookkeeping re: trainable parameters.
-            assert circuit.qtape.trainable_params == {0, 1, 2, 3, 4}
+        assert circuit.qtape.trainable_params == {1, 2, 3, 4}
 
         assert [i.name for i in circuit.qtape.operations] == ["RX", "Rot", "PhaseShift"]
-
-        if diff_method == "finite-diff":
-            assert np.all(circuit.qtape.get_parameters() == [p[2], p[0], -p[2], p[1] + p[2]])
-        elif diff_method == "backprop":
-            # In backprop mode, all parameters are returned.
-            assert np.all(circuit.qtape.get_parameters() == [a, p[2], p[0], -p[2], p[1] + p[2]])
+        assert np.all(circuit.qtape.get_parameters() == [p[2], p[0], -p[2], p[1] + p[2]])
 
         expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
             np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
diff --git a/tests/interfaces/test_qnode_tf.py b/tests/interfaces/test_qnode_tf.py
index c674c3b5c..ff728a49a 100644
--- a/tests/interfaces/test_qnode_tf.py
+++ b/tests/interfaces/test_qnode_tf.py
@@ -447,20 +447,9 @@ class TestQNode:
         with tf.GradientTape() as tape:
             res = circuit(a, p)
 
-        if diff_method == "finite-diff":
-            assert circuit.qtape.trainable_params == {1, 2, 3, 4}
-        elif diff_method == "backprop":
-            # For a backprop device, no interface wrapping is performed, and JacobianTape.jacobian()
-            # is never called. As a result, JacobianTape.trainable_params is never set --- the ML
-            # framework uses its own backprop logic and its own bookkeeping re: trainable parameters.
-            assert circuit.qtape.trainable_params == {0, 1, 2, 3, 4}
-
+        assert circuit.qtape.trainable_params == {1, 2, 3, 4}
         assert [i.name for i in circuit.qtape.operations] == ["RX", "Rot", "PhaseShift"]
-
-        if diff_method == "finite-diff":
-            assert np.all(circuit.qtape.get_parameters() == [p[2], p[0], -p[2], p[1] + p[2]])
-        elif diff_method == "backprop":
-            assert np.all(circuit.qtape.get_parameters() == [a, p[2], p[0], -p[2], p[1] + p[2]])
+        assert np.all(circuit.qtape.get_parameters() == [p[2], p[0], -p[2], p[1] + p[2]])
 
         expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * (
             tf.cos(p[2]) * tf.sin(p[1]) + tf.cos(p[0]) * tf.cos(p[1]) * tf.sin(p[2])
diff --git a/tests/interfaces/test_qnode_torch.py b/tests/interfaces/test_qnode_torch.py
index f0ca8ce7a..7a0bc1ec9 100644
--- a/tests/interfaces/test_qnode_torch.py
+++ b/tests/interfaces/test_qnode_torch.py
@@ -458,20 +458,9 @@ class TestQNode:
 
         res = circuit(a, p)
 
-        if diff_method == "finite-diff":
-            assert circuit.qtape.trainable_params == {1, 2, 3, 4}
-        elif diff_method == "backprop":
-            # For a backprop device, no interface wrapping is performed, and JacobianTape.jacobian()
-            # is never called. As a result, JacobianTape.trainable_params is never set --- the ML
-            # framework uses its own backprop logic and its own bookkeeping re: trainable parameters.
-            assert circuit.qtape.trainable_params == {0, 1, 2, 3, 4}
-
+        assert circuit.qtape.trainable_params == {1, 2, 3, 4}
         assert [i.name for i in circuit.qtape.operations] == ["RX", "Rot", "PhaseShift"]
-
-        if diff_method == "finite-diff":
-            assert np.all(circuit.qtape.get_parameters() == [p[2], p[0], -p[2], p[1] + p[2]])
-        elif diff_method == "backprop":
-            assert np.all(circuit.qtape.get_parameters() == [a, p[2], p[0], -p[2], p[1] + p[2]])
+        assert np.all(circuit.qtape.get_parameters() == [p[2], p[0], -p[2], p[1] + p[2]])
 
         expected = np.cos(a) * np.cos(p_val[1]) * np.sin(p_val[0]) + np.sin(a) * (
             np.cos(p_val[2]) * np.sin(p_val[1])
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 4304b3c69..bf8286508 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -1610,8 +1610,7 @@ class TestGetTrainable:
         values = [[0.1, 0.2], np.tensor(0.1, requires_grad=True), np.tensor([0.5, 0.2])]
         cost_fn(values)
 
-        # Currently, we assume *all* objects are trainable by default in Autograd
-        assert res == {0, 1, 2}
+        assert res == {1, 2}
 
     def test_autograd_unwrapping_backward(self):
         """Test that the trainability indices of a sequence of Autograd arrays
diff --git a/tests/tape/test_unwrap.py b/tests/tape/test_unwrap.py
index d6a8a947a..f6cf1f33f 100644
--- a/tests/tape/test_unwrap.py
+++ b/tests/tape/test_unwrap.py
@@ -99,7 +99,7 @@ def test_unwrap_autograd():
         params = tape.get_parameters(trainable_only=False)
         assert all(isinstance(i, float) for i in params)
         assert np.allclose(params, [0.1, 0.2, 0.5, 0.3])
-        assert tape.trainable_params == {0, 2, 3}
+        assert tape.trainable_params == {0, 3}
 
     # outside the context, the original parameters have been restored.
     assert tape.get_parameters(trainable_only=False) == p
diff --git a/tests/transforms/test_classical_jacobian.py b/tests/transforms/test_classical_jacobian.py
index e4286764b..4548ef4f4 100644
--- a/tests/transforms/test_classical_jacobian.py
+++ b/tests/transforms/test_classical_jacobian.py
@@ -111,11 +111,12 @@ class_jacs = [
 ]
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
-def test_autograd_without_argnum(circuit, args, expected_jac):
+def test_autograd_without_argnum(circuit, args, expected_jac, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=None`` and Autograd."""
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="autograd")
+    qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method)
     jac = classical_jacobian(qnode)(*args)
 
     # NOTE: We use stacking to replicate qml.jacobian behaviour for scalar-only inputs
@@ -131,27 +132,28 @@ def test_autograd_without_argnum(circuit, args, expected_jac):
             assert np.allclose(_jac, _expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
-def test_jax_without_argnum(circuit, args, expected_jac):
+def test_jax_without_argnum(circuit, args, expected_jac, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=None`` and JAX."""
     # Do not need the package but skip if JAX device not available
     pytest.importorskip("jax")
     # JAX behaviour: argnum=None yields only the Jacobian with respect to the first arg.
     expected_jac = expected_jac[0]
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="jax")
+    qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method)
     jac = classical_jacobian(qnode)(*args)
-
     assert np.allclose(jac, expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
-def test_tf_without_argnum(circuit, args, expected_jac):
+def test_tf_without_argnum(circuit, args, expected_jac, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=None`` and Tensorflow."""
     tf = pytest.importorskip("tensorflow")
     args = tuple((tf.Variable(arg, dtype=tf.double) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="tf")
+    qnode = qml.QNode(circuit, dev, interface="tf", diff_method=diff_method)
     jac = classical_jacobian(qnode)(*args)
 
     assert len(jac) == len(expected_jac)
@@ -159,13 +161,14 @@ def test_tf_without_argnum(circuit, args, expected_jac):
         assert np.allclose(_jac, _expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
-def test_torch_without_argnum(circuit, args, expected_jac):
+def test_torch_without_argnum(circuit, args, expected_jac, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=None`` and Torch."""
     torch = pytest.importorskip("torch")
     args = tuple((torch.tensor(arg) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="torch")
+    qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method)
     jac = classical_jacobian(qnode)(*args)
 
     assert len(jac) == len(expected_jac)
@@ -176,55 +179,59 @@ def test_torch_without_argnum(circuit, args, expected_jac):
 scalar_argnum = [0, 1, 0, 1, 0, 1]
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
 )
-def test_autograd_with_scalar_argnum(circuit, args, expected_jac, argnum):
+def test_autograd_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Autograd."""
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="autograd")
+    qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = expected_jac[argnum]
     assert np.allclose(jac, expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
 )
-def test_jax_with_scalar_argnum(circuit, args, expected_jac, argnum):
+def test_jax_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and JAX."""
     # Do not need the package but skip if JAX device not available
     pytest.importorskip("jax")
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="jax")
+    qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = expected_jac[argnum]
     assert np.allclose(jac, expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
 )
-def test_tf_with_scalar_argnum(circuit, args, expected_jac, argnum):
+def test_tf_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and TensorFlow."""
     tf = pytest.importorskip("tensorflow")
     args = tuple((tf.Variable(arg, dtype=tf.double) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="tf")
+    qnode = qml.QNode(circuit, dev, interface="tf", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = expected_jac[argnum]
     assert np.allclose(jac, expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, scalar_argnum)
 )
-def test_torch_with_scalar_argnum(circuit, args, expected_jac, argnum):
+def test_torch_with_scalar_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Torch."""
     torch = pytest.importorskip("torch")
     args = tuple((torch.tensor(arg) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="torch")
+    qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = expected_jac[argnum]
     assert np.allclose(jac, expected_jac)
@@ -233,58 +240,62 @@ def test_torch_with_scalar_argnum(circuit, args, expected_jac, argnum):
 single_list_argnum = [[0], [1], [0], [1], [0], [2]]
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
 )
-def test_autograd_with_single_list_argnum(circuit, args, expected_jac, argnum):
+def test_autograd_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and Autograd."""
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="autograd")
+    qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = (expected_jac[argnum[0]],)
     assert len(jac) == 1
     assert np.allclose(jac[0], expected_jac[0])
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
 )
-def test_jax_with_single_list_argnum(circuit, args, expected_jac, argnum):
+def test_jax_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and JAX."""
     # Do not need the package but skip if JAX device not available
     pytest.importorskip("jax")
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="jax")
+    qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = (expected_jac[argnum[0]],)
     assert len(jac) == 1
     assert np.allclose(jac[0], expected_jac[0])
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
 )
-def test_tf_with_single_list_argnum(circuit, args, expected_jac, argnum):
+def test_tf_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and TensorFlow."""
     tf = pytest.importorskip("tensorflow")
     args = tuple((tf.Variable(arg, dtype=tf.double) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="tf")
+    qnode = qml.QNode(circuit, dev, interface="tf", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = (expected_jac[argnum[0]],)
     assert len(jac) == 1
     assert np.allclose(jac[0], expected_jac[0])
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, single_list_argnum)
 )
-def test_torch_with_single_list_argnum(circuit, args, expected_jac, argnum):
+def test_torch_with_single_list_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` of length 1 and Torch."""
     torch = pytest.importorskip("torch")
     args = tuple((torch.tensor(arg) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="torch")
+    qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = (expected_jac[argnum[0]],)
     assert len(jac) == 1
@@ -294,13 +305,14 @@ def test_torch_with_single_list_argnum(circuit, args, expected_jac, argnum):
 sequence_argnum = [[0], [0, 1], (0,), [0, 1], (0, 1), {0, 2}]
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
 )
-def test_autograd_with_sequence_argnum(circuit, args, expected_jac, argnum):
+def test_autograd_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and Autograd."""
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="autograd")
+    qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = tuple((expected_jac[num] for num in argnum))
     assert len(jac) == len(expected_jac)
@@ -308,15 +320,16 @@ def test_autograd_with_sequence_argnum(circuit, args, expected_jac, argnum):
         assert np.allclose(_jac, _expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
 )
-def test_jax_with_sequence_argnum(circuit, args, expected_jac, argnum):
+def test_jax_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and JAX."""
     # Do not need the package but skip if JAX device not available
     pytest.importorskip("jax")
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="jax")
+    qnode = qml.QNode(circuit, dev, interface="jax", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = tuple((expected_jac[num] for num in argnum))
     assert len(jac) == len(expected_jac)
@@ -324,15 +337,16 @@ def test_jax_with_sequence_argnum(circuit, args, expected_jac, argnum):
         assert np.allclose(_jac, _expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
 )
-def test_tf_with_sequence_argnum(circuit, args, expected_jac, argnum):
+def test_tf_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and TensorFlow."""
     tf = pytest.importorskip("tensorflow")
     args = tuple((tf.Variable(arg, dtype=tf.double) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="tf")
+    qnode = qml.QNode(circuit, dev, interface="tf", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = tuple((expected_jac[num] for num in argnum))
     assert len(jac) == len(expected_jac)
@@ -340,15 +354,16 @@ def test_tf_with_sequence_argnum(circuit, args, expected_jac, argnum):
         assert np.allclose(_jac, _expected_jac)
 
 
+@pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift"])
 @pytest.mark.parametrize(
     "circuit, args, expected_jac, argnum", zip(circuits, args, class_jacs, sequence_argnum)
 )
-def test_torch_with_sequence_argnum(circuit, args, expected_jac, argnum):
+def test_torch_with_sequence_argnum(circuit, args, expected_jac, argnum, diff_method):
     r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and Torch."""
     torch = pytest.importorskip("torch")
     args = tuple((torch.tensor(arg) for arg in args))
     dev = qml.device("default.qubit", wires=2)
-    qnode = qml.QNode(circuit, dev, interface="torch")
+    qnode = qml.QNode(circuit, dev, interface="torch", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = tuple((expected_jac[num] for num in argnum))
     assert len(jac) == len(expected_jac)
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index 56eac8987..76ff4370d 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -15,7 +15,7 @@
 Unit tests for the metric tensor transform.
 """
 import pytest
-import numpy as np
+from pennylane import numpy as np
 from scipy.linalg import block_diag
 
 import pennylane as qml
@@ -29,14 +29,13 @@ class TestMetricTensor:
     def test_rot_decomposition(self, diff_method):
         """Test that the rotation gate is correctly decomposed"""
         dev = qml.device("default.qubit", wires=1)
+        params = np.array([1.0, 2.0, 3.0], requires_grad=True)
 
-        def circuit(weights):
-            qml.Rot(weights[0], weights[1], weights[2], wires=0)
-            return qml.expval(qml.PauliX(0))
+        with qml.tape.QuantumTape() as circuit:
+            qml.Rot(params[0], params[1], params[2], wires=0)
+            qml.expval(qml.PauliX(0))
 
-        circuit = qml.QNode(circuit, dev, diff_method=diff_method)
-        params = np.array([1.0, 2.0, 3.0])
-        tapes = qml.metric_tensor(circuit, only_construct=True)(params)
+        tapes, _ = qml.metric_tensor(circuit)
         assert len(tapes) == 3
 
         # first parameter subcircuit
@@ -58,9 +57,6 @@ class TestMetricTensor:
         assert tapes[2].operations[0].data == [1]
         assert tapes[2].operations[1].data == [2]
 
-        result = qml.metric_tensor(circuit)(params)
-        assert result.shape == (3, 3)
-
     @pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop"])
     def test_multirz_decomposition(self, diff_method):
         """Test that the MultiRZ gate is correctly decomposed"""
@@ -92,39 +88,34 @@ class TestMetricTensor:
 
         circuit = qml.QNode(circuit, dev, diff_method=diff_method)
         params = [0.1]
-        result = qml.metric_tensor(circuit)(*params)
+        result = qml.metric_tensor(circuit, hybrid=False)(*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"""
-        dev = qml.device("default.qubit", wires=1)
-
-        def circuit(a):
-            qml.RX(a, wires=0)
-            return qml.expval(qml.PauliX(0))
-
-        circuit = qml.QNode(circuit, dev)
-
         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"):
-                circuit.metric_tensor(1.0, only_construct=True)
+                qml.metric_tensor(tape)
 
     def test_construct_subcircuit(self):
         """Test correct subcircuits constructed"""
         dev = qml.device("default.qubit", wires=2)
 
-        def circuit(a, b, c):
-            qml.RX(a, wires=0)
-            qml.RY(b, wires=0)
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(np.array(1.0, requires_grad=True), wires=0)
+            qml.RY(np.array(1.0, requires_grad=True), wires=0)
             qml.CNOT(wires=[0, 1])
-            qml.PhaseShift(c, wires=1)
+            qml.PhaseShift(np.array(1.0, requires_grad=True), wires=1)
             return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
 
-        circuit = qml.QNode(circuit, dev)
-        tapes = circuit.metric_tensor(1.0, 1.0, 1.0, only_construct=True)
+        tapes, _ = qml.metric_tensor(tape)
         assert len(tapes) == 3
 
         # first parameter subcircuit
@@ -151,8 +142,9 @@ class TestMetricTensor:
         """Test correct subcircuits constructed
         when a layer structure exists"""
         dev = qml.device("default.qubit", wires=3)
+        params = np.ones([8])
 
-        def circuit(params):
+        with qml.tape.QuantumTape() as tape:
             # section 1
             qml.RX(params[0], wires=0)
             # section 2
@@ -173,10 +165,7 @@ class TestMetricTensor:
             qml.CNOT(wires=[1, 2])
             return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1)), qml.expval(qml.PauliX(2))
 
-        circuit = qml.QNode(circuit, dev)
-
-        params = np.ones([8])
-        tapes = circuit.metric_tensor(params, only_construct=True)
+        tapes, _ = qml.metric_tensor(tape)
 
         # this circuit should split into 4 independent
         # sections or layers when constructing subcircuits
@@ -243,7 +232,7 @@ class TestMetricTensor:
         c = -0.432
 
         # evaluate metric tensor
-        g = circuit.metric_tensor(a, b, c)
+        g = qml.metric_tensor(circuit)(a, b, c)
 
         # check that the metric tensor is correct
         expected = (
@@ -254,6 +243,41 @@ class TestMetricTensor:
         )
         assert np.allclose(g, np.diag(expected), atol=tol, rtol=0)
 
+    def test_evaluate_diag_metric_tensor_classical_processing(self, tol):
+        """Test that a diagonal metric tensor evaluates correctly
+        when the QNode includes classical processing."""
+        dev = qml.device("default.qubit", wires=2)
+
+        def circuit(a, b):
+            # The classical processing function is
+            #     f: ([a0, a1], b) -> (a1, a0, b)
+            # So the classical Jacobians will be a permutation matrix and an identity matrix:
+            #     classical_jacobian(circuit)(a, b) == ([[0, 1], [1, 0]], [[1]])
+            qml.RX(a[1], wires=0)
+            qml.RY(a[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PhaseShift(b, wires=1)
+            return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
+
+        circuit = qml.QNode(circuit, dev)
+
+        a = np.array([0.432, 0.1])
+        b = 0.12
+
+        # evaluate metric tensor
+        g = qml.metric_tensor(circuit)(a, b)
+        assert isinstance(g, tuple)
+        assert len(g) == 2
+        assert g[0].shape == (len(a), len(a))
+        assert g[1].shape == tuple()
+
+        # 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)
+
+        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)
+
     @pytest.fixture(params=["parameter-shift", "backprop"])
     def sample_circuit(self, request):
         """Sample variational circuit fixture used in the
@@ -300,7 +324,7 @@ class TestMetricTensor:
         params = [-0.282203, 0.145554, 0.331624, -0.163907, 0.57662, 0.081272]
         x, y, z, h, g, f = params
 
-        G = circuit.metric_tensor(*params)
+        G = qml.metric_tensor(circuit)(*params)
 
         # ============================================
         # Test block diag metric tensor of first layer is correct.
@@ -443,7 +467,7 @@ class TestMetricTensor:
         params = [-0.282203, 0.145554, 0.331624, -0.163907, 0.57662, 0.081272]
         x, y, z, h, g, f = params
 
-        G = circuit.metric_tensor(*params, diag_approx=True)
+        G = qml.metric_tensor(circuit, diag_approx=True)(*params)
 
         # ============================================
         # Test block diag metric tensor of first layer is correct.
@@ -553,7 +577,7 @@ class TestDifferentiability:
         def cost(weights):
             return qml.metric_tensor(circuit)(weights)[2, 2]
 
-        weights = np.array([0.432, 0.12, -0.432])
+        weights = np.array([0.432, 0.12, -0.432], requires_grad=True)
         a, b, c = weights
 
         grad = qml.grad(cost)(weights)
@@ -647,3 +671,60 @@ class TestDifferentiability:
             [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)
+
+
+class TestDeprecatedQNodeMethod:
+    """The QNode.metric_tensor method has been deprecated.
+    These tests ensure it still works, but raises a deprecation
+    warning. These tests can be deleted when the method is removed."""
+
+    def test_warning(self, tol):
+        """Test that a warning is emitted"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(a, b, c):
+            qml.RX(a, wires=0)
+            qml.RY(b, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PhaseShift(c, wires=1)
+            return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
+
+        a = 0.432
+        b = 0.12
+        c = -0.432
+
+        # evaluate metric tensor
+        with pytest.warns(UserWarning, match="has been deprecated"):
+            g = circuit.metric_tensor(a, b, c)
+
+        # check that the metric tensor is correct
+        expected = (
+            np.array(
+                [1, np.cos(a) ** 2, (3 - 2 * np.cos(a) ** 2 * np.cos(2 * b) - np.cos(2 * a)) / 4]
+            )
+            / 4
+        )
+        assert np.allclose(g, np.diag(expected), atol=tol, rtol=0)
+
+    def test_tapes_returned(self, tol):
+        """Test that a warning is emitted"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(a, b, c):
+            qml.RX(a, wires=0)
+            qml.RY(b, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PhaseShift(c, wires=1)
+            return qml.expval(qml.PauliX(0)), qml.expval(qml.PauliX(1))
+
+        a = 0.432
+        b = 0.12
+        c = -0.432
+
+        # evaluate metric tensor
+        with pytest.warns(UserWarning, match="has been deprecated"):
+            tapes, fn = circuit.metric_tensor(a, b, c, only_construct=True)
+
+        assert len(tapes) == 3