Repository: PennyLaneAI/pennylane
Pull request #3587 (2023-01-02 13:59:51)

Add `argnum` argument to `metric_tensor`
**Context:**
The [metric_tensor](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/transforms/metric_tensor.py) function is constructed by differentiating a quantum tape with respect to all trainable parameters and computing inner products between the respective quantum states. When one is interested in only a subset of all metric-tensor elements one might want to restrict the set of trainable parameters to save computational resources. The `argnum` argument allows the caller to do this.

**Description of the Change:**
The `argnum` argument must be either an integer or a list of integers or `None`. The default is `None` which reproduces the previous behavior, i.e. derivatives with respect to _all_ trainable parameters are computed. When an integer or a list of integers is given only the corresponding tape parameters are taken into account. All other metric tensor entries are set to 0.

Example:
```Python
dev = qml.device("default.qubit", wires=3)

def circuit(weights):
    qml.RX(weights[0], wires=0)
    qml.RY(weights[1], wires=0)
    qml.CNOT(wires=[0, 1])
    qml.RZ(weights[2], wires=1)
    qml.RZ(weights[3], wires=0)

weights = np.array([0.1, 0.2, 0.3, 0.5], requires_grad=True)

with qml.tape.QuantumTape() as tape:
    circuit(weights)

tapes, proc_fn = qml.metric_tensor(tape)
res = qml.execute(tapes, dev, None)
print(proc_fn(res))

# [[ 0.25        0.         -0.04966733 -0.04966733]
#  [ 0.          0.24750832  0.02433865  0.02433865]
#  [-0.04966733  0.02433865  0.01226071  0.01226071]
#  [-0.04966733  0.02433865  0.01226071  0.01226071]]
```
We can restrict ourselves to only the parameters 2 and 3.
```Python
tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
res = qml.execute(tapes, dev, None)
print(proc_fn(res))

# [[0.         0.         0.         0.        ]
#  [0.         0.         0.         0.        ]
#  [0.         0.         0.01226071 0.01226071]
#  [0.         0.         0.01226071 0.01226071]]
```
Notably, the first evaluation required 6 quantum tapes, while the second one only requires one tape, as both remaining parameters are contained in one layer.

**Benefits:**
When only a reduced version of the metric tensor is needed, this argument allows a convenient way to reduce required quantum resources.

**Possible Drawbacks:**
The indices of a QNode argument might not be identical to the corresponding tape parameters. However, the `argnum` argument has to specify the parameters by tape indices. This might lead to confusing behavior for a user that tries to compute the metric tensor of a QNode. To mitigate this to some extend, a warning is given when an index is specified that does not refer to any trainable tape parameter.

**Related GitHub Issues:**
This PR addresses #1880.

**Note**
I am not sure the test I added is extensive enough. The argument works with all three `approx` methods. Since the `approx=None` case covers the entire code used to compute the block-diagonal and diagonal approximation, I don't think one has to test these individually.

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3587?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3587](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3587?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (f2779f6) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/880517228462ba01af35d94204e55fc9c47a7cac?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (8805172) will **decrease** coverage by `0.01%`.
> The diff coverage is `95.12%`.

> :exclamation: Current head f2779f6 differs from pull request most recent head 423847c. Consider uploading reports for the commit 423847c to get more accurate results

```diff
@@            Coverage Diff             @@
##           master    #3587      +/-   ##
==========================================
- Coverage   99.87%   99.85%   -0.02%     
==========================================
  Files         314      314              
  Lines       28019    28077      +58     
==========================================
+ Hits        27983    28037      +54     
- Misses         36       40       +4     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3587?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3587/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWV0cmljX3RlbnNvci5weQ==) | `98.35% <95.12%> (-1.65%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

frederikwilde: Thanks for this thorough review, @dwierichs ! I addressed many comments already. The two main things that remain are

1. The potential performance issues arising from the way the metric-tensor blocks are built in `_metric_tensor_cov_matrix`. I'll take a look at this soon.
2. The dummy lists in `obs_list` and `coeffs_list` that would avoid a lot of awkward `for` looping. The issue here is, as I stated in a direct response to your comment, that the `metric_tensor_tapes` list would then also need some sort of "dummy tape" corresponding to the dummy lists. Let me know if you have any idea how to do this.
dwierichs: Great @frederikwilde, thank you for the edits!
This is in nice shape already! :100: 

> The issue here is, as I stated in a direct response to your comment, that the metric_tensor_tapes list would then also need some sort of "dummy tape" corresponding to the dummy lists. 

See my response to the comment regarding this. I think there is a nice middle ground. A single counter that compensates for "missing" tapes definitely is fine, but there shouldn't be too many, in my opinion.

> Also CodeFactor is complaining about functions having more than 5 arguments. I can try to refactor some things to resolve this. What do you think?

If there is a nice part to make into a smaller function and/or reduce the number of arguments, that would be nice. However, sometimes it is not easily doable or actually decreases readability. Maybe you can give it a try!

Could I also ask you to add some small test for the interfaces (see e.g. `test_tf_diag`, `test_tf` functions) that uses `argnum`? Thanks! :)
frederikwilde: Thanks for all the comments! I think we are getting pretty close. Just wanted to mention briefly, that I will get back to this soon. The next two days are just quite busy.
frederikwilde: > Could I also ask you to add some small test for the interfaces (see e.g. test_tf_diag, test_tf functions) that uses argnum? Thanks! :)

Isn't this already covered by this [parametrization](https://github.com/frederikwilde/pennylane/blob/add-argnum-to-metric_tensor/tests/transforms/test_metric_tensor.py#L582) of the test function?
dwierichs: Hi @frederikwilde,
do you require more information/help/review with this PR to get it merge-ready? :)
frederikwilde: @dwierichs Thanks for the ping. I am on vacation currently. I tried to finish the PR before I left, but got caught up in other work and a paper review. I will get back to this PR at latest after Feb. 11.
rmoyard: Hi @frederikwilde! If you don't have time to do it this week could you pls allow me to work on the branch directly (setting in the PR that you can change I think)
frederikwilde: Sure, but you already committed directly to my branch, no? I allowed maintainers to edit, that's the only setting I have for this. Lmk if it doesn't work.
frederikwilde: Thank you for all the feedback. The example in the doc string is still missing. I could do this now. What do you think? There are also two typos in line 91 and 153.
rmoyard: Thanks @frederikwilde I let you finish it and then we will merge it, thank you!
frederikwilde: Done. Feel free to adapt it as you like. I think one could add a note about the amount of tapes one saves by setting `argnum` to different values (see initial post in this PR), but maybe this is also a bit too much for the doc string and also not exactly a usage detail.

Thanks @rmoyard, for resolving my awkward index tracking, I couldn't really come up with an elegant solution.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 064870533..f0864f8a4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -542,6 +542,11 @@
   `qml.gradients.SUPPORTED_GRADIENT_KWARGS`.
   [(#3526)](https://github.com/PennyLaneAI/pennylane/pull/3526)
 
+* Added `argnum` argument to `metric_tensor`. By passing a sequence of indices referring to trainable tape parameters,
+  the metric tensor is only computed with respect to these parameters. This reduces the number of tapes that have to
+  be run.
+  [(#3587)](https://github.com/PennyLaneAI/pennylane/pull/3587)
+
 * Moved `qml.utils.sparse_hamiltonian` function to `~.Hamiltonian.sparse_matrix` method.
   [(#3585)](https://github.com/PennyLaneAI/pennylane/pull/3585)
 
@@ -816,4 +821,5 @@ Mudit Pandey
 Borja Requena
 Matthew Silverman
 Antal Száva
+Frederik Wilde
 David Wierichs
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index 954f4569f..a17d928a2 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -25,7 +25,9 @@ from pennylane.circuit_graph import LayerData
 from .batch_transform import batch_transform
 
 
-def expand_fn(tape, approx=None, allow_nonunitary=True, aux_wire=None, device_wires=None):
+def expand_fn(
+    tape, argnum=None, 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,too-many-arguments
     if not allow_nonunitary and approx is None:  # pragma: no cover
@@ -34,7 +36,9 @@ def expand_fn(tape, approx=None, allow_nonunitary=True, aux_wire=None, device_wi
 
 
 @functools.partial(batch_transform, expand_fn=expand_fn)
-def metric_tensor(tape, approx=None, allow_nonunitary=True, aux_wire=None, device_wires=None):
+def metric_tensor(
+    tape, argnum=None, approx=None, allow_nonunitary=True, aux_wire=None, device_wires=None
+):  # pylint: disable=too-many-arguments
     r"""Returns a function that computes the metric tensor of a given QNode or quantum tape.
 
     The metric tensor convention we employ here has the following form:
@@ -58,6 +62,12 @@ def metric_tensor(tape, approx=None, allow_nonunitary=True, aux_wire=None, devic
 
     Args:
         tape (pennylane.QNode or .QuantumTape): quantum tape or QNode to find the metric tensor of
+        argnum (int or Sequence[int] or None): Trainable tape-parameter indices with respect to which
+
+            the metric tensor is computed. If `argnum=None`, the metric tensor with respect to all
+            trainable parameters is returned. Excluding tape-parameter indices from this list reduces
+            the computational cost and the corresponding metric-tensor elements will be set to 0.
+
         approx (str): Which approximation of the metric tensor to compute.
 
             - If ``None``, the full metric tensor is computed
@@ -243,6 +253,41 @@ def metric_tensor(tape, approx=None, allow_nonunitary=True, aux_wire=None, devic
 
         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.
+
+        .. warning::
+
+            The ``argnum`` argument can be used to restrict the parameters which are taken into account
+            for computing the metric tensor.
+            When the metric tensor of a QNode is computed, the ordering of the parameters has to be
+            specified as they appear in the corresponding QuantumTape.
+
+        **Example**
+
+        Consider the following QNode in which parameters are used out of order:
+
+        .. code-block:: python
+
+            >>> dev = qml.device("default.qubit", wires=3)
+            >>> @qml.qnode(dev, interface="autograd")
+            >>> def circuit(weights):  # , extra_weight):
+            ...     qml.RX(weights[1], wires=0)
+            ...     qml.RY(weights[0], wires=0)
+            ...     qml.CNOT(wires=[0, 1])
+            ...     qml.RZ(weights[2], wires=1)
+            ...     qml.RZ(weights[3], wires=0)
+            ...     return qml.expval(qml.PauliZ(0))
+
+            >>> mt = qml.metric_tensor(circuit, argnum=(0, 2, 3))(weights)
+            >>> print(mt)
+            [[ 0.          0.          0.          0.        ]
+            [ 0.          0.25       -0.02495835 -0.02495835]
+            [ 0.         -0.02495835  0.01226071  0.01226071]
+            [ 0.         -0.02495835  0.01226071  0.01226071]]
+
+        Because the 0-th element of ``weights`` appears second in the QNode and therefore in the
+        underlying tape, it is the 1st tape parameter.
+        By setting ``argnum = (0, 2, 3)`` we exclude the 0-th element of ``weights`` from the computation
+        of the metric tensor and not the 1st element, as one might expect.
     """
     if not tape.trainable_params:
         warnings.warn(
@@ -252,14 +297,27 @@ def metric_tensor(tape, approx=None, allow_nonunitary=True, aux_wire=None, devic
         )
         return [], lambda _: ()
 
+    if argnum is None:
+        argnum = tape.trainable_params
+    elif isinstance(argnum, int):
+        argnum = [argnum]
+    if any(i not in tape.trainable_params for i in argnum):
+        raise ValueError(
+            "Some parameters specified in argnum are not in the "
+            f"trainable parameters {tape.trainable_params} of the tape "
+            "and will be ignored. This may be caused by attempting to "
+            "differentiate with respect to parameters that are not marked "
+            "as trainable."
+        )
+
     # pylint: disable=too-many-arguments
     if approx in {"diag", "block-diag"}:
         # Only require covariance matrix based transform
         diag_approx = approx == "diag"
-        return _metric_tensor_cov_matrix(tape, diag_approx)[:2]
+        return _metric_tensor_cov_matrix(tape, argnum, diag_approx)[:2]
 
     if approx is None:
-        return _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires)
+        return _metric_tensor_hadamard(tape, argnum, allow_nonunitary, aux_wire, device_wires)
 
     raise ValueError(
         f"Unknown value {approx} for keyword argument approx. "
@@ -379,12 +437,14 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
     return wrapper
 
 
-def _metric_tensor_cov_matrix(tape, diag_approx):
+def _metric_tensor_cov_matrix(tape, argnum, diag_approx):  # pylint: disable=too-many-statements
     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
+        argnum (list[int]): Trainable tape-parameter indices with respect to which the metric tensor
+            is computed.
         diag_approx (bool): if True, use the diagonal approximation. If ``False`` , a
             block-diagonal approximation of the metric tensor is computed.
     Returns:
@@ -396,6 +456,12 @@ def _metric_tensor_cov_matrix(tape, diag_approx):
             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
+        list[list[bool]]: Each inner list corresponds to one tape and therefore also one parametrized
+            layer and its elements determine whether a trainable parameter in that layer is
+            contained in ``argnum``.
+        list[None, int]: Id list representing the layer for each parameter.
+        list[None, int]: Id list representing the observables for each parameter.
+
 
     This method assumes the ``generator`` of all parametrized operations with respect to
     which the tensor is computed to be Hermitian. This is the case for unitary single-parameter
@@ -408,27 +474,52 @@ def _metric_tensor_cov_matrix(tape, diag_approx):
     obs_list = []
     coeffs_list = []
     params_list = []
+    in_argnum_list = []
+    layers_ids = []
+    obs_ids = []
 
+    i = 0
     for queue, curr_ops, param_idx, _ in graph.iterate_parametrized_layers():
         params_list.append(param_idx)
-        coeffs_list.append([])
-        obs_list.append([])
+        in_argnum_list.append([p in argnum for p in param_idx])
+
+        if not any(in_argnum_list[-1]):
+            layers_ids.extend([None] * len(in_argnum_list[-1]))
+            obs_ids.extend([None] * len(in_argnum_list[-1]))
+            # no tape needs to be created for this block
+            continue
+
+        layer_coeffs, layer_obs = [], []
 
         # for each operation in the layer, get the generator
-        for op in curr_ops:
-            obs, s = qml.generator(op)
-            obs_list[-1].append(obs)
-            coeffs_list[-1].append(s)
+        j = 0
+        for p, op in zip(param_idx, curr_ops):
+            layers_ids.append(i)
+            if p in argnum:
+                obs, s = qml.generator(op)
+                layer_obs.append(obs)
+                layer_coeffs.append(s)
+                obs_ids.append(j)
+                j = j + 1
+            else:
+                obs_ids.append(None)
+        i = i + 1
+
+        coeffs_list.append(layer_coeffs)
+        obs_list.append(layer_obs)
 
         # Create a quantum tape with all operations
         # prior to the parametrized layer, and the rotations
         # to measure in the basis of the parametrized layer generators.
         with qml.queuing.AnnotatedQueue() as layer_q:
             for op in queue:
+                # TODO: Maybe there are gates that do not affect the
+                # generators of interest and thus need not be applied.
                 qml.apply(op)
 
-            for o in obs_list[-1]:
-                o.diagonalizing_gates()
+            for o, param_in_argnum in zip(layer_obs, in_argnum_list[-1]):
+                if param_in_argnum:
+                    o.diagonalizing_gates()
 
             qml.probs(wires=tape.wires)
 
@@ -437,18 +528,45 @@ def _metric_tensor_cov_matrix(tape, diag_approx):
 
     def processing_fn(probs):
         gs = []
-
-        for prob, obs, coeffs in zip(probs, obs_list, coeffs_list):
-            # calculate the covariance matrix of this layer
-            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)
+        probs_idx = 0
+
+        for params_in_argnum in in_argnum_list:
+            if not any(params_in_argnum):
+                # there is no tape and no probs associated to this layer
+                dim = len(params_in_argnum)
+                gs.append(qml.math.zeros((dim, dim)))
+                continue
+
+            coeffs = coeffs_list[probs_idx]
+            obs = obs_list[probs_idx]
+            p = probs[probs_idx]
+
+            scale = qml.math.convert_like(qml.math.outer(coeffs, coeffs), p)
+            scale = qml.math.cast_like(scale, p)
+            g = scale * qml.math.cov_matrix(p, obs, wires=tape.wires, diag_approx=diag_approx)
+            for i, in_argnum in enumerate(params_in_argnum):
+                # fill in rows and columns of zeros where a parameter was not in argnum
+                if not in_argnum:
+                    dim = g.shape[0]
+                    g = qml.math.concatenate((g[:i], qml.math.zeros((1, dim)), g[i:]))
+                    g = qml.math.concatenate(
+                        (g[:, :i], qml.math.zeros((dim + 1, 1)), g[:, i:]), axis=1
+                    )
             gs.append(g)
+            probs_idx += 1
 
         # create the block diagonal metric tensor
         return qml.math.block_diag(gs)
 
-    return metric_tensor_tapes, processing_fn, obs_list, coeffs_list
+    return (
+        metric_tensor_tapes,
+        processing_fn,
+        obs_list,
+        coeffs_list,
+        in_argnum_list,
+        layers_ids,
+        obs_ids,
+    )
 
 
 @functools.lru_cache()
@@ -553,13 +671,17 @@ def _get_first_term_tapes(layer_i, layer_j, allow_nonunitary, aux_wire):
     return tapes, ids
 
 
-def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
+def _metric_tensor_hadamard(
+    tape, argnum, allow_nonunitary, aux_wire, device_wires
+):  # pylint: disable=too-many-statements
     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
+        argnum (list[int]): Trainable tape-parameter indices with respect to which the metric tensor
+            is computed.
         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.
@@ -576,21 +698,38 @@ def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
     """
     # 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)
+    (
+        diag_tapes,
+        diag_proc_fn,
+        obs_list,
+        coeffs_list,
+        in_argnum_list,
+        layer_ids,
+        obs_ids,
+    ) = _metric_tensor_cov_matrix(tape, argnum, 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 = [
-        LayerData(
-            layer.pre_ops,
-            layer.ops,
-            tuple(par_idx_to_trainable_idx[idx] for idx in layer[2]),
-            layer.post_ops,
-        )
-        for layer in graph.iterate_parametrized_layers()
-    ]
+    layers = []
+
+    for layer, in_argnum in zip(graph.iterate_parametrized_layers(), in_argnum_list):
+        if not any(in_argnum):
+            # no tapes need to be constructed for this layer
+            continue
+
+        pre_ops, ops, param_idx, post_ops = layer
+        new_ops = []
+        new_param_idx = []
+
+        for o, idx, param_in_argnum in zip(ops, param_idx, in_argnum):
+            if param_in_argnum:
+                new_ops.append(o)
+                new_param_idx.append(par_idx_to_trainable_idx[idx])
+
+        layers.append(LayerData(pre_ops, new_ops, new_param_idx, post_ops))
+
     if len(layers) <= 1:
         return diag_tapes, diag_proc_fn
 
@@ -613,7 +752,12 @@ def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
     # 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])
+    blocks = []
+    for in_argnum in in_argnum_list:
+        d = len(in_argnum)
+        blocks.append(qml.math.ones((d, d)))
+    mask = 1 - qml.math.block_diag(blocks)
+
     # Required for slicing in processing_fn
     num_diag_tapes = len(diag_tapes)
 
@@ -642,14 +786,15 @@ def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
 
         # 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:
+
+        for i, (layer_i, obs_i) in enumerate(zip(layer_ids, obs_ids)):
+            if layer_i is not None and obs_i is not None:
+                prob = diag_res[layer_i]
+                o = obs_list[layer_i][obs_i]
                 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
+                expvals = qml.math.scatter_element_add(expvals, (i,), qml.math.dot(l, p))
 
         # Construct <\partial_i\psi|\psi><\psi|\partial_j\psi> and mask it
         second_term = qml.math.tensordot(expvals, expvals, axes=0) * mask
@@ -658,8 +803,16 @@ def _metric_tensor_hadamard(tape, allow_nonunitary, aux_wire, device_wires):
         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])
+        coeffs_gen = (c for c in qml.math.hstack(coeffs_list))
+        # flattened coeffs_list but also with 0s where parameters are not in argnum
+        extended_coeffs_list = [
+            next(coeffs_gen) if param_in_argnum else 0.0
+            for param_in_argnum in qml.math.hstack(in_argnum_list)
+        ]
+
+        scale = qml.math.convert_like(
+            qml.math.tensordot(extended_coeffs_list, extended_coeffs_list, axes=0), results[0]
+        )
         off_diag_mt = scale * off_diag_mt
 
         # Combine block-diagonal and off block-diagonal


Test patch:
diff --git a/tests/returntypes/transforms/test_metric_tensor_new.py b/tests/returntypes/transforms/test_metric_tensor_new.py
index 34027923e..55de62782 100644
--- a/tests/returntypes/transforms/test_metric_tensor_new.py
+++ b/tests/returntypes/transforms/test_metric_tensor_new.py
@@ -581,6 +581,267 @@ class TestMetricTensor:
         G_expected = block_diag(G1, G3, G2)
         assert qml.math.allclose(G, G_expected, atol=tol, rtol=0)
 
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("interface", ["autograd"])
+    def test_argnum_metric_tensor_autograd(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        dev = qml.device("default.qubit.autograd", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = np.array([0.1, 0.2, 0.3, 0.5], requires_grad=True)
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, np.ndarray)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, np.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, np.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("interface", ["tf"])
+    def test_argnum_metric_tensor_tf(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.tf", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = tf.Variable([0.1, 0.2, 0.3, 0.5])
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, tf.Tensor)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, tf.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, tf.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("interface", ["torch"])
+    def test_argnum_metric_tensor_torch(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        import torch
+
+        dev = qml.device("default.qubit.torch", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = torch.tensor([0.1, 0.2, 0.3, 0.5], requires_grad=True)
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, torch.Tensor)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, torch.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, torch.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_argnum_metric_tensor_jax(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        import jax
+
+        dev = qml.device("default.qubit.jax", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = jax.numpy.array([0.1, 0.2, 0.3, 0.5])
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, jax.numpy.ndarray)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, jax.numpy.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, jax.numpy.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    def test_argnum_metric_tensor_errors(self):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        dev = qml.device("default.qubit", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = np.array([0.1, 0.2, 0.3, 0.5], requires_grad=True)
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        error_msg = (
+            "Some parameters specified in argnum are not in the "
+            "trainable parameters \[0, 1, 2, 3\] of the tape "
+            "and will be ignored. This may be caused by attempting to "
+            "differentiate with respect to parameters that are not marked "
+            "as trainable."
+        )
+        with pytest.raises(ValueError, match=error_msg):
+            tapes, proc_fn = qml.metric_tensor(tape, argnum=4)
+
     def test_multi_qubit_gates(self):
         """Test that a tape with Ising gates has the correct metric tensor tapes."""
 
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index b5cbd0779..c5c548c53 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -580,6 +580,267 @@ class TestMetricTensor:
         G_expected = block_diag(G1, G3, G2)
         assert qml.math.allclose(G, G_expected, atol=tol, rtol=0)
 
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("interface", ["autograd"])
+    def test_argnum_metric_tensor_autograd(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        dev = qml.device("default.qubit.autograd", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = np.array([0.1, 0.2, 0.3, 0.5], requires_grad=True)
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, np.ndarray)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, np.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, np.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("interface", ["tf"])
+    def test_argnum_metric_tensor_tf(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.tf", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = tf.Variable([0.1, 0.2, 0.3, 0.5])
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, tf.Tensor)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, tf.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, tf.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("interface", ["torch"])
+    def test_argnum_metric_tensor_torch(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        import torch
+
+        dev = qml.device("default.qubit.torch", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = torch.tensor([0.1, 0.2, 0.3, 0.5], requires_grad=True)
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, torch.Tensor)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, torch.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, torch.Tensor)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_argnum_metric_tensor_jax(self, tol, interface):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        import jax
+
+        dev = qml.device("default.qubit.jax", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = jax.numpy.array([0.1, 0.2, 0.3, 0.5])
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        tapes, proc_fn = qml.metric_tensor(tape)
+        res = qml.execute(tapes, dev, None)
+        mt = proc_fn(res)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt013 = proc_fn(res)
+        assert isinstance(mt013, jax.numpy.ndarray)
+
+        assert len(tapes) == 6
+        assert mt.shape == mt013.shape
+        assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0)
+        assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3))
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt23 = proc_fn(res)
+        assert isinstance(mt23, jax.numpy.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt23.shape
+        assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0)
+
+        tapes, proc_fn = qml.metric_tensor(tape, argnum=0)
+        res = qml.execute(tapes, dev, None, interface=interface)
+        mt0 = proc_fn(res)
+        assert isinstance(mt0, jax.numpy.ndarray)
+
+        assert len(tapes) == 1
+        assert mt.shape == mt0.shape
+        assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0)
+        assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0)
+
+    def test_argnum_metric_tensor_errors(self):
+        """Test that argnum successfully reduces the number of tapes and gives
+        the desired outcome."""
+        dev = qml.device("default.qubit", wires=3)
+
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[2], wires=1)
+            qml.RZ(weights[3], wires=0)
+
+        weights = np.array([0.1, 0.2, 0.3, 0.5], requires_grad=True)
+
+        with qml.tape.QuantumTape() as tape:
+            circuit(weights)
+
+        error_msg = (
+            "Some parameters specified in argnum are not in the "
+            "trainable parameters \[0, 1, 2, 3\] of the tape "
+            "and will be ignored. This may be caused by attempting to "
+            "differentiate with respect to parameters that are not marked "
+            "as trainable."
+        )
+        with pytest.raises(ValueError, match=error_msg):
+            tapes, proc_fn = qml.metric_tensor(tape, argnum=4)
+
     def test_multi_qubit_gates(self):
         """Test that a tape with Ising gates has the correct metric tensor tapes."""