Repository: PennyLaneAI/pennylane
Pull request #1645 (2021-09-11 07:45:17)

Add `argnum` keyword argument and tests to `classical_jacobian`
**Context:**
The new keyword argument `argnum` allows to compute the classical preprocessing Jacobian of a QNode only for selected QNode arguments. 

**Description of the Change:**
If not provided with `argnum`, i.e. `argnum=None`, the output type/content depends on the interface.
If provided with an integer `argnum`, the corresponding Jacobian array is returned as single output.
If provided with an iterable `argnum`, the ouput is a tuple of arrays.
See the docstring for details on the behaviour, which is non-trivial for the default case `argnum=None`.

**Benefits:**
1. Added tests for `classical_jacobian`. Previously, it was only tested implicitly via `transforms/test_metric_tensor.py`.
2. Flexibility of `classical_jacobian`. In particular, we will use this in a new version of `qml.fourier.spectrum` and for `qml.fourier.reconstruct` (two upcoming PRs).
3. `classical_jacobian` previously did not work for QNodes that take multiple arguments of different shapes. This is now possible.

**Possible Drawbacks:**
None

**Related GitHub Issues:**
#1628  (Note that the previously WIP PR used to contain this PR)


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1645?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1645](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1645?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e15dfb7) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/1bb2bf19e31bc9882dd7012f7189530005a21212?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1bb2bf1) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1645/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1645?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1645   +/-   ##
=======================================
  Coverage   99.14%   99.14%           
=======================================
  Files         196      196           
  Lines       14248    14271   +23     
=======================================
+ Hits        14126    14149   +23     
  Misses        122      122           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1645?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1645/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `100.00% <100.00%> (ø)` | |

------

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

dwierichs: @josh146 I believe this is ready for review :)

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index acee6e233..61738e66c 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,42 @@
 
 <h3>New features since last release</h3>
 
+* The transform for the Jacobian of the classical preprocessing within a QNode,
+  `qml.transforms.classical_jacobian`, now takes a keyword argument `argnum` to specify
+  the QNode argument indices with respect to which the Jacobian is computed.
+  [(#1645)](https://github.com/PennyLaneAI/pennylane/pull/1645)
+
+  An example for the usage of ``argnum`` is
+
+  ```python
+  @qml.qnode(dev)
+  def circuit(x, y, z):
+      qml.RX(qml.math.sin(x), wires=0)
+      qml.CNOT(wires=[0, 1])
+      qml.RY(y ** 2, wires=1)
+      qml.RZ(1 / z, wires=1)
+      return qml.expval(qml.PauliZ(0))
+
+  jac_fn = qml.transforms.classical_jacobian(circuit, argnum=[1, 2])
+  ```
+
+  The Jacobian can then be computed at specified parameters.
+
+  ```pycon
+  >>> x, y, z = np.array([0.1, -2.5, 0.71])
+  >>> jac_fn(x, y, z)
+  (array([-0., -5., -0.]), array([-0.        , -0.        , -1.98373339]))
+  ```
+
+  The returned arrays are the derivatives of the three parametrized gates in the circuit
+  with respect to `y` and `z` respectively.
+
+  There also are explicit tests for `classical_jacobian` now, which previously was tested
+  implicitly via its use in the `metric_tensor` transform.
+
+  For more usage details, please see the
+  [classical Jacobian docstring](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.transforms.classical_jacobian.html).
+
 <h3>Improvements</h3>
 
 * ``qml.circuit_drawer.CircuitDrawer`` can accept a string for the ``charset`` keyword, instead of a ``CharSet`` object.
@@ -17,6 +53,8 @@
 
 This release contains contributions from (in alphabetical order):
 
+Christina Lee, David Wierichs
+
 # Release 0.18.0 (Current release)
 
 <h3>New features since last release</h3>
@@ -73,6 +111,7 @@ This release contains contributions from (in alphabetical order):
   [(#1225)](https://github.com/PennyLaneAI/pennylane/pull/1360)
   [(#1598)](https://github.com/PennyLaneAI/pennylane/pull/1598)
 
+
 * The ability to define *batch* transforms has been added via the new
   `@qml.batch_transform` decorator.
   [(#1493)](https://github.com/PennyLaneAI/pennylane/pull/1493)
diff --git a/pennylane/transforms/classical_jacobian.py b/pennylane/transforms/classical_jacobian.py
index b19e80860..e377fc9e3 100644
--- a/pennylane/transforms/classical_jacobian.py
+++ b/pennylane/transforms/classical_jacobian.py
@@ -12,13 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-Contains the classical Jacobian transform
+Contains the classical Jacobian transform.
 """
 # pylint: disable=import-outside-toplevel
 import pennylane as qml
+from pennylane import numpy as np
 
 
-def classical_jacobian(qnode):
+def classical_jacobian(qnode, argnum=None):
     r"""Returns a function to extract the Jacobian
     matrix of the classical part of a QNode.
 
@@ -27,11 +28,13 @@ def classical_jacobian(qnode):
 
     Args:
         qnode (.QNode): QNode to compute the (classical) Jacobian of
+        argnum (int or Sequence[int]): indices of QNode arguments with respect to which
+            the (classical) Jacobian is computed
 
     Returns:
         function: Function which accepts the same arguments as the QNode.
         When called, this function will return the Jacobian of the QNode
-        gate arguments with respect to the QNode arguments.
+        gate arguments with respect to the QNode arguments indexed by ``argnum``.
 
     **Example**
 
@@ -40,21 +43,25 @@ def classical_jacobian(qnode):
     >>> @qml.qnode(dev)
     ... def circuit(weights):
     ...     qml.RX(weights[0], wires=0)
-    ...     qml.RY(weights[0], wires=1)
-    ...     qml.RZ(weights[2] ** 2, wires=1)
-    ...     return qml.expval(qml.PauliZ(0))
+    ...     qml.RY(0.2 * weights[0], wires=1)
+    ...     qml.RY(2.5, wires=0)
+    ...     qml.RZ(weights[1] ** 2, wires=1)
+    ...     qml.RX(weights[2], wires=1)
+    ...     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
     We can use this transform to extract the relationship :math:`f: \mathbb{R}^n \rightarrow
     \mathbb{R}^m` between the input QNode arguments :math:`w` and the gate arguments :math:`g`, for
     a given value of the QNode arguments:
 
     >>> cjac_fn = qml.transforms.classical_jacobian(circuit)
-    >>> weights = np.array([1., 1., 1.], requires_grad=True)
+    >>> weights = np.array([1., 1., 0.6], requires_grad=True)
     >>> cjac = cjac_fn(weights)
     >>> print(cjac)
-    [[1. 0. 0.]
-     [1. 0. 0.]
-     [0. 0. 2.]]
+    [[1.  0.  0. ]
+     [0.2 0.  0. ]
+     [0.  0.  0. ]
+     [0.  1.2 0. ]
+     [0.  0.  1. ]]
 
     The returned Jacobian has rows corresponding to gate arguments, and columns
     corresponding to QNode arguments; that is,
@@ -65,39 +72,129 @@ def classical_jacobian(qnode):
 
     - The zeroth element of ``weights`` is repeated on the first two gates generated by the QNode.
 
-    - The second column consisting of all zeros indicates that the generated quantum circuit does
-      not depend on the first element of ``weights``.
+    - The third row consisting of all zeros indicates that the third gate ``RY(2.5)`` does not
+      depend on the ``weights``.
+
+    - The quadratic dependence of the fourth gate argument yields :math:`2\cdot 0.6=1.2`.
+
+    .. note::
+
+        The QNode is constructed during this operation.
+
+    For a QNode with multiple QNode arguments, the arguments with respect to which the
+    Jacobian is computed can be controlled with the ``argnum`` keyword argument.
+    The output for ``argnum=None`` depends on the backend:
+
+    .. list-table:: Output format of ``classical_jacobian``
+       :widths: 25 25 25 25
+       :header-rows: 1
+
+       * - Interface
+         - ``argnum=None``
+         - ``type(argnum)=int``
+         - ``argnum=Sequence[int]``
+       * - ``'autograd'``
+         - ``tuple(arrays)`` [1]
+         - ``array``
+         - ``tuple(array)``
+       * - ``'jax'``
+         - ``array``
+         - ``array``
+         - ``tuple(array)``
+       * - ``'tf'``
+         - ``tuple(arrays)``
+         - ``array``
+         - ``tuple(array)``
+       * - ``'torch'``
+         - ``tuple(arrays)``
+         - ``array``
+         - ``tuple(array)``
+
+    [1] If all QNode argument are scalars, the tuple is unpacked and the one-dimensional Jacobian
+    arrays are stacked into one ``array``. If there only is one QNode argument, the tuple is
+    unpacked as well. Both is due to the behaviour of ``qml.jacobian``.
+
+    **Example with ``argnum``**
+
+    >>> @qml.qnode(dev)
+    ... def circuit(x, y, z):
+    ...     qml.RX(qml.math.sin(x), wires=0)
+    ...     qml.CNOT(wires=[0, 1])
+    ...     qml.RY(y ** 2, wires=1)
+    ...     qml.RZ(1 / z, wires=1)
+    ...     return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+    >>> jac_fn = qml.transforms.classical_jacobian(circuit, argnum=[1, 2])
+    >>> x, y, z = np.array([0.1, -2.5, 0.71])
+    >>> jac_fn(x, y, z)
+    (array([-0., -5., -0.]), array([-0.        , -0.        , -1.98373339]))
+
+    Only the Jacobians with respect to the arguments ``x`` and ``y`` were computed, and
+    returned as a tuple of ``arrays``.
+
     """
 
     def classical_preprocessing(*args, **kwargs):
-        """Returns the trainable gate parameters for
-        a given QNode input"""
+        """Returns the trainable gate parameters for a given QNode input."""
         qnode.construct(args, kwargs)
         return qml.math.stack(qnode.qtape.get_parameters())
 
     if qnode.interface == "autograd":
-        return qml.jacobian(classical_preprocessing)
+
+        def _jacobian(*args, **kwargs):
+            if argnum is None:
+                jac = qml.jacobian(classical_preprocessing)(*args, **kwargs)
+            elif np.isscalar(argnum):
+                jac = qml.jacobian(classical_preprocessing, argnum=argnum)(*args, **kwargs)
+            else:
+                jac = tuple(
+                    (
+                        qml.jacobian(classical_preprocessing, argnum=i)(*args, **kwargs)
+                        for i in argnum
+                    )
+                )
+            return jac
+
+        return _jacobian
 
     if qnode.interface == "torch":
         import torch
 
         def _jacobian(*args, **kwargs):  # pylint: disable=unused-argument
-            return torch.autograd.functional.jacobian(classical_preprocessing, args)
+            jac = torch.autograd.functional.jacobian(classical_preprocessing, args)
+            if argnum is not None:
+                if np.isscalar(argnum):
+                    jac = jac[argnum]
+                else:
+                    jac = tuple((jac[idx] for idx in argnum))
+            return jac
 
         return _jacobian
 
     if qnode.interface == "jax":
         import jax
 
-        return jax.jacobian(classical_preprocessing)
+        argnum = 0 if argnum is None else argnum
+
+        def _jacobian(*args, **kwargs):
+            return jax.jacobian(classical_preprocessing, argnums=argnum)(*args, **kwargs)
+
+        return _jacobian
 
     if qnode.interface == "tf":
         import tensorflow as tf
 
         def _jacobian(*args, **kwargs):
+            if np.isscalar(argnum):
+                sub_args = args[argnum]
+            elif argnum is None:
+                sub_args = args
+            else:
+                sub_args = tuple((args[i] for i in argnum))
+
             with tf.GradientTape() as tape:
                 gate_params = classical_preprocessing(*args, **kwargs)
 
-            return tape.jacobian(gate_params, args)
+            jac = tape.jacobian(gate_params, sub_args)
+            return jac
 
         return _jacobian


Test patch:
diff --git a/tests/transforms/test_classical_jacobian.py b/tests/transforms/test_classical_jacobian.py
new file mode 100644
index 000000000..e4286764b
--- /dev/null
+++ b/tests/transforms/test_classical_jacobian.py
@@ -0,0 +1,356 @@
+# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the qml.transforms.classical_jacobian function."""
+import pytest
+import numpy as np
+
+import pennylane as qml
+from pennylane.transforms.classical_jacobian import classical_jacobian
+
+a = -2.1
+b = 0.71
+x = np.array([0.3, 2.3, 0.1])
+y = np.array([[1.0, 2.0], [4.0, 5.0]])
+z = np.array([2.1, -0.3, 0.62, 0.89])
+
+
+def circuit_0(a):
+    [qml.RX(a, wires=0) for i in range(4)]
+    return qml.expval(qml.PauliZ(0))
+
+
+def circuit_1(a, b):
+    qml.RX(qml.math.sin(a), wires=0)
+    qml.RZ(a / 3, wires=0)
+    qml.CNOT(wires=[0, 1])
+    qml.RY(b ** 2, wires=1)
+    qml.RZ(1 / b, wires=1)
+    return qml.expval(qml.PauliZ(0))
+
+
+def circuit_2(x):
+    [qml.RX(x[i], wires=0) for i in range(3)]
+    return qml.expval(qml.PauliZ(0))
+
+
+def circuit_3(x, y):
+    [qml.RX(x[i], wires=0) for i in range(3)]
+    for i in range(2):
+        [qml.RY(y[i, j], wires=1) for j in range(2)]
+    return qml.expval(qml.PauliZ(0))
+
+
+perm_3 = ([2, 0, 1], [1, 2, 0, 3])
+
+
+def circuit_4(x, y):
+    for i in perm_3[0]:
+        qml.RX(x[i], wires=0)
+    for j in perm_3[1]:
+        qml.RY(y[j // 2, j % 2], wires=1)
+    return qml.expval(qml.PauliZ(0))
+
+
+def circuit_5(x, y, z):
+    [qml.RX(x[i], wires=0) for i in range(3)]
+    qml.RZ(y[0, 1] * y[1, 0], wires=1)
+    qml.RY(z[0] + 0.2 * z[1] ** 2, wires=1)
+    return qml.expval(qml.PauliZ(0))
+
+
+circuits = [circuit_0, circuit_1, circuit_2, circuit_3, circuit_4, circuit_5]
+args = [(a,), (a, b), (x,), (x, y), (x, y), (x, y, z)]
+interfaces = ["jax", "autograd", "tf", "torch"]
+
+class_jacs = [
+    (np.ones(4),),
+    (
+        np.array(
+            [
+                np.cos(a),
+                1 / 3,
+                0.0,
+                0.0,
+            ]
+        ),
+        np.array([0.0, 0.0, 2 * b, -1 / (b ** 2)]),
+    ),
+    (np.eye(len(x)),),
+    (
+        np.vstack([np.eye(len(x)), np.zeros((4, 3))]),
+        np.vstack([np.zeros((3,) + y.shape), np.eye(np.prod(y.shape)).reshape(-1, *y.shape)]),
+    ),
+    (
+        np.vstack([np.eye(len(x)), np.zeros((4, 3))])[perm_3[0] + [3, 4, 5, 6]],
+        np.vstack(
+            [np.zeros((3,) + y.shape), np.eye(np.prod(y.shape))[perm_3[1]].reshape(-1, *y.shape)]
+        ),
+    ),
+    (
+        np.vstack([np.eye(len(x)), np.zeros((2, 3))]),
+        np.vstack(
+            [
+                np.zeros((3,) + y.shape),
+                np.array([[[0.0, y[1, 0]], [y[0, 1], 0.0]]]),
+                np.zeros((1,) + y.shape),
+            ]
+        ),
+        np.vstack([np.zeros((4, 4)), np.array([1, 0.4 * z[1], 0.0, 0.0])]),
+    ),
+]
+
+
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+def test_autograd_without_argnum(circuit, args, expected_jac):
+    r"""Test ``classical_jacobian`` with ``argnum=None`` and Autograd."""
+    dev = qml.device("default.qubit", wires=2)
+    qnode = qml.QNode(circuit, dev, interface="autograd")
+    jac = classical_jacobian(qnode)(*args)
+
+    # NOTE: We use stacking to replicate qml.jacobian behaviour for scalar-only inputs
+    if all((np.isscalar(arg) for arg in args)):
+        expected_jac = qml.math.stack(expected_jac).T
+        assert np.allclose(jac, expected_jac)
+    else:
+        # For a single argument, the Jacobian is unpacked
+        if len(args) == 1:
+            expected_jac = expected_jac[0]
+        assert len(jac) == len(expected_jac)
+        for _jac, _expected_jac in zip(jac, expected_jac):
+            assert np.allclose(_jac, _expected_jac)
+
+
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+def test_jax_without_argnum(circuit, args, expected_jac):
+    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")
+    jac = classical_jacobian(qnode)(*args)
+
+    assert np.allclose(jac, expected_jac)
+
+
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+def test_tf_without_argnum(circuit, args, expected_jac):
+    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")
+    jac = classical_jacobian(qnode)(*args)
+
+    assert len(jac) == len(expected_jac)
+    for _jac, _expected_jac in zip(jac, expected_jac):
+        assert np.allclose(_jac, _expected_jac)
+
+
+@pytest.mark.parametrize("circuit, args, expected_jac", zip(circuits, args, class_jacs))
+def test_torch_without_argnum(circuit, args, expected_jac):
+    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")
+    jac = classical_jacobian(qnode)(*args)
+
+    assert len(jac) == len(expected_jac)
+    for _jac, _expected_jac in zip(jac, expected_jac):
+        assert np.allclose(_jac, _expected_jac)
+
+
+scalar_argnum = [0, 1, 0, 1, 0, 1]
+
+
+@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):
+    r"""Test ``classical_jacobian`` with ``argnum=<int>`` and Autograd."""
+    dev = qml.device("default.qubit", wires=2)
+    qnode = qml.QNode(circuit, dev, interface="autograd")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = expected_jac[argnum]
+    assert np.allclose(jac, expected_jac)
+
+
+@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):
+    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")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = expected_jac[argnum]
+    assert np.allclose(jac, expected_jac)
+
+
+@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):
+    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")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = expected_jac[argnum]
+    assert np.allclose(jac, expected_jac)
+
+
+@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):
+    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")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = expected_jac[argnum]
+    assert np.allclose(jac, expected_jac)
+
+
+single_list_argnum = [[0], [1], [0], [1], [0], [2]]
+
+
+@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):
+    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")
+    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(
+    "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):
+    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")
+    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(
+    "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):
+    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")
+    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(
+    "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):
+    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")
+    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])
+
+
+sequence_argnum = [[0], [0, 1], (0,), [0, 1], (0, 1), {0, 2}]
+
+
+@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):
+    r"""Test ``classical_jacobian`` with ``argnum=Sequence[int]`` and Autograd."""
+    dev = qml.device("default.qubit", wires=2)
+    qnode = qml.QNode(circuit, dev, interface="autograd")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = tuple((expected_jac[num] for num in argnum))
+    assert len(jac) == len(expected_jac)
+    for _jac, _expected_jac in zip(jac, expected_jac):
+        assert np.allclose(_jac, _expected_jac)
+
+
+@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):
+    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")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = tuple((expected_jac[num] for num in argnum))
+    assert len(jac) == len(expected_jac)
+    for _jac, _expected_jac in zip(jac, expected_jac):
+        assert np.allclose(_jac, _expected_jac)
+
+
+@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):
+    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")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = tuple((expected_jac[num] for num in argnum))
+    assert len(jac) == len(expected_jac)
+    for _jac, _expected_jac in zip(jac, expected_jac):
+        assert np.allclose(_jac, _expected_jac)
+
+
+@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):
+    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")
+    jac = classical_jacobian(qnode, argnum=argnum)(*args)
+    expected_jac = tuple((expected_jac[num] for num in argnum))
+    assert len(jac) == len(expected_jac)
+    for _jac, _expected_jac in zip(jac, expected_jac):
+        assert np.allclose(_jac, _expected_jac)