Repository: PennyLaneAI/pennylane (python)

Issue:
Allow tape.jacobian to estimate the jacobian by computing the gradients of a subset of parameters chosen stochastically
Currently, in PennyLane, the `tape.jacobian` method will loop over all parameters in the circuit to construct the Jacobian. If utilizing the parameter-shift rule method of differentiation, for example, this will execute the variational circuit `2*num_parameter` times in order to compute the gradient.

However, a recent result by [Ryan Sweke et al. (2019)](https://arxiv.org/abs/1910.01155) showed that you can randomly sample the parameters to shift at each optimization step. Even in the extreme case where only one parameter is shifted per step, gradient descent convergence remains guaranteed.

This would be a nice feature to add to PennyLane. This could be done as follows:

* Allow a new keyword argument option to be passed to `JacobianTape.jacobian`, for example `k_params` (this is just a placeholder name, final name to be determined).

* Then, modify the internal for loop over all parameters,

  ```python
  for trainable_idx, param_method in enumerate(diff_methods):
  ```

  so that gradient rules are generated by randomly sampling only the specified number of parameters, *if the new keyword argument is provided*:

  ```python
  k_params = options.get("k_params", None)

  params_and_methods = enumerate(diff_methods)
  if stochastic_params is not None:
      params_and_methods = set(random.choices(list(enumerate(diff_methods)), k=k_params))
  
  for trainable_idx, param_method in params_and_methods:
  ```
  


Pull request description:
Allow specifying a subset of parameters to estimate the jacobian by passing the `argnum` kwarg
**Context:**
#958 

**Description of the Change:**

Adds the `argnum` keyword argument to the `JacobianTape.jacobian` method inspired by the `argnum` kwarg of the `qml.grad` function. Using this keyword argument the trainable parameters to use for estimating the jacobian can be specified.

The logic is organized into a separate private method to allow easy unit testing and to prepare for further future extensions.

*Note:* as this is an addition to the `JacobianTape`, this option is only applicable for the differentiation methods that depend on such an object and its `jacobian` method: ``"parameter-shift"``, ``"finite-diff"`` and ``"reversible"``.

**Benefits:**

Jacobians can now be estimated by only using a subset of the trainable parameters.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
Closes #958

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 2981f8e00..f2f7b4f34 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -2,6 +2,43 @@
 
 <h3>New features since last release</h3>
 
+* The ``argnum`` keyword argument can now be specified for a QNode to define a
+  subset of trainable parameters used to estimate the Jacobian.
+  [(#1371)](https://github.com/PennyLaneAI/pennylane/pull/1371)
+
+  For example, consider two trainable parameters and a quantum function:
+
+  ```python
+  dev = qml.device("default.qubit", wires=2)
+
+  x = np.array(0.543, requires_grad=True)
+  y = np.array(-0.654, requires_grad=True)
+
+  def circuit(x,y):
+      qml.RX(x, wires=[0])
+      qml.RY(y, wires=[1])
+      qml.CNOT(wires=[0, 1])
+      return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+  ```
+
+  When computing the gradient of the QNode, we can specify the trainable
+  parameters to consider by passing the ``argnum`` keyword argument:
+
+  ```pycon
+  >>> qnode1 = qml.QNode(circuit, dev, diff_method="parameter-shift", argnum=[0,1])
+  >>> print(qml.grad(qnode1)(x,y))
+  (array(0.31434679), array(0.67949903))
+  ```
+
+  Specifying a proper subset of the trainable parameters will estimate the
+  Jacobian:
+
+  ```pycon
+  >>> qnode2 = qml.QNode(circuit, dev, diff_method="parameter-shift", argnum=[0])
+  >>> print(qml.grad(qnode2)(x,y))
+  (array(0.31434679), array(0.))
+  ```
+
 * The `quantum_monte_carlo` transform has been added, allowing an input circuit to be transformed
   into the full quantum Monte Carlo algorithm.
   [(#1316)](https://github.com/PennyLaneAI/pennylane/pull/1316)
@@ -72,15 +109,15 @@
 * Added validation for noise channel parameters. Invalid noise parameters now
   raise a `ValueError`. [(#1357)](https://github.com/PennyLaneAI/pennylane/pull/1357)
 
-* PennyLane now has a `fourier` module, which hosts a [growing library 
-  of methods](https://pennylane.readthedocs.io/en/stable/code/qml_fourier.html) 
-  that help with investigating the Fourier representation of functions 
+* PennyLane now has a `fourier` module, which hosts a [growing library
+  of methods](https://pennylane.readthedocs.io/en/stable/code/qml_fourier.html)
+  that help with investigating the Fourier representation of functions
   implemented by quantum circuits.
   [(#1160)](https://github.com/PennyLaneAI/pennylane/pull/1160)
-  
-  For example, one can plot distributions over Fourier series coefficients like 
+
+  For example, one can plot distributions over Fourier series coefficients like
   this one:
-  
+
   <img src="https://pennylane.readthedocs.io/en/latest/_static/fourier.png" width=70%/>
 
 * It is now possible [(1291)](https://github.com/PennyLaneAI/pennylane/pull/1291)
@@ -315,7 +352,7 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
   `adjoint_cache=False` during QNode creation.
   [(#1341)](https://github.com/PennyLaneAI/pennylane/pull/1341)
 
-* The `Operator` (and by inheritance, the `Operation` and `Observable` class and their children) 
+* The `Operator` (and by inheritance, the `Operation` and `Observable` class and their children)
   now have an `id` attribute, which can mark an operator in a circuit, for example to
   identify it on the tape by a tape transform.
   [(#1377)](https://github.com/PennyLaneAI/pennylane/pull/1377)
@@ -422,10 +459,10 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 
 <h3>Bug fixes</h3>
 
-* Fixes bug where `shots=None` was not reset when changing shots temporarily in a QNode call 
+* Fixes bug where `shots=None` was not reset when changing shots temporarily in a QNode call
   like `circuit(0.1, shots=3)`.
   [(#1392)](https://github.com/XanaduAI/pennylane/pull/1392)
-  
+
 * Fixes floating point errors with `diff_method="finite-diff"` and `order=1` when parameters are `float32`.
 [(#1381)](https://github.com/PennyLaneAI/pennylane/pull/1381)
 
@@ -462,7 +499,7 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 
 <h3>Documentation</h3>
 
-* Fix typo in the documentation of qml.templates.layers.StronglyEntanglingLayers 
+* Fix typo in the documentation of qml.templates.layers.StronglyEntanglingLayers
   [(#1367)](https://github.com/PennyLaneAI/pennylane/pull/1367)
 
 * Fixed typo on TensorFlow interface documentation [(#1312)](https://github.com/PennyLaneAI/pennylane/pull/1312)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 78861cc64..dd766a09c 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -126,6 +126,11 @@ class QNode:
         adjoint_cache=True (bool): for TensorFlow and PyTorch interfaces and adjoint differentiation,
             this indicates whether to save the device state after the forward pass.  Doing so saves a
             forward execution. Device state automatically reused with autograd and JAX interfaces.
+        argnum=None (int, list(int), None): Which argument(s) to compute the Jacobian
+            with respect to. When there are fewer parameters specified than the
+            total number of trainable parameters, the jacobian is being estimated. Note
+            that this option is only applicable for the following differentiation methods:
+            ``"parameter-shift"``, ``"finite-diff"`` and ``"reversible"``.
 
     **Example**
 
@@ -954,6 +959,11 @@ def qnode(
         adjoint_cache=True (bool): for TensorFlow and PyTorch interfaces and adjoint differentiation,
             this indicates whether to save the device state after the forward pass.  Doing so saves a
             forward execution. Device state automatically reused with autograd and JAX interfaces.
+        argnum=None (int, list(int), None): Which argument(s) to compute the Jacobian
+            with respect to. When there are fewer parameters specified than the
+            total number of trainable parameters, the jacobian is being estimated. Note
+            that this option is only applicable for the following differentiation methods:
+            ``"parameter-shift"``, ``"finite-diff"`` and ``"reversible"``.
 
     **Example**
 
diff --git a/pennylane/tape/jacobian_tape.py b/pennylane/tape/jacobian_tape.py
index dc1e640e8..9843b18b0 100644
--- a/pennylane/tape/jacobian_tape.py
+++ b/pennylane/tape/jacobian_tape.py
@@ -18,10 +18,11 @@ to the ``QuantumTape`` class.
 # pylint: disable=too-many-branches
 
 import itertools
+import warnings
+
 import numpy as np
 
 import pennylane as qml
-
 from pennylane.operation import State
 from pennylane.tape import QuantumTape
 
@@ -398,6 +399,47 @@ class JacobianTape(QuantumTape):
         """
         raise NotImplementedError
 
+    def _choose_params_with_methods(self, diff_methods, argnum):
+        """Chooses the trainable parameters to use for computing the Jacobian
+        by returning a map of their indices and differentiation methods.
+
+        When there are fewer parameters specified than the total number of
+        trainable parameters, the Jacobian is estimated by using the parameters
+        specified using the ``argnum`` keyword argument.
+
+        Args:
+            diff_methods (list): the ordered list of differentiation methods
+                for each parameter
+            argnum (int, list(int), None): Indices for which argument(s) to
+                compute the Jacobian with respect to.
+
+        Returns:
+            enumerate or list: map of the trainable parameter indices and
+            differentiation methods
+        """
+        if argnum is None:
+            return enumerate(diff_methods)
+
+        if isinstance(argnum, int):
+            argnum = [argnum]
+
+        if not all(ind in self.trainable_params for ind in argnum):
+            raise ValueError(
+                "Incorrect trainable parameters were specified for the argnum argument."
+            )
+
+        num_params = len(argnum)
+
+        if num_params == 0:
+            warnings.warn(
+                "No trainable parameters were specified for computing the Jacobian.",
+                UserWarning,
+            )
+            return []
+
+        diff_methods_to_use = map(diff_methods.__getitem__, argnum)
+        return zip(argnum, diff_methods_to_use)
+
     def jacobian(self, device, params=None, **options):
         r"""Compute the Jacobian of the parametrized quantum circuit recorded by the quantum tape.
 
@@ -446,6 +488,9 @@ class JacobianTape(QuantumTape):
             order=1 (int): The order of the finite difference method to use. ``1`` corresponds
                 to forward finite differences, ``2`` to centered finite differences.
             shift=pi/2 (float): the size of the shift for two-term parameter-shift gradient computations
+            argnum=None (int, list(int), None): Which argument(s) to compute the Jacobian
+                with respect to. When there are fewer parameters specified than the
+                total number of trainable parameters, the jacobian is being estimated.
 
         Returns:
             array[float]: 2-dimensional array of shape ``(tape.output_dim, tape.num_params)``
@@ -554,7 +599,11 @@ class JacobianTape(QuantumTape):
         processing_fns = []
         nonzero_grad_idx = []
 
-        for trainable_idx, param_method in enumerate(diff_methods):
+        argnum = options.get("argnum", None)
+
+        params_with_methods = self._choose_params_with_methods(diff_methods, argnum)
+
+        for trainable_idx, param_method in params_with_methods:
             if param_method == "0":
                 continue
 


Test patch:
diff --git a/tests/tape/test_jacobian_tape.py b/tests/tape/test_jacobian_tape.py
index fd9aa7a90..528d25039 100644
--- a/tests/tape/test_jacobian_tape.py
+++ b/tests/tape/test_jacobian_tape.py
@@ -431,6 +431,57 @@ class TestJacobian:
         assert np.allclose(j1, [exp, 0])
         assert np.allclose(j2, [0, exp])
 
+    @pytest.mark.parametrize(
+        "diff_methods", [["A", "0", "F"], ["A", "A", "A"], ["F", "A", "A", "0", "0"]]
+    )
+    @pytest.mark.parametrize("argnum", [None, 0, [0, 1], [0, 1, 2], [2, 0], [1, 0], [0, 0, 0]])
+    def test_choose_params_and_methods(self, diff_methods, argnum):
+        """Test that the _choose_params_and_methods helper method returns
+        expected results"""
+        tape = JacobianTape()
+        tape._trainable_params = list(range(len(diff_methods)))
+        res = list(tape._choose_params_with_methods(diff_methods, argnum))
+
+        num_all_params = len(diff_methods)
+
+        assert all(k in range(num_all_params) for k, _ in res)
+        assert all(v in diff_methods for _, v in res)
+
+        if argnum is None:
+            num_params = num_all_params
+        elif isinstance(argnum, int):
+            num_params = 1
+        else:
+            num_params = len(argnum)
+
+        assert len(res) == num_params
+
+    @pytest.mark.parametrize("argnum", [1, 2, 3, -1])
+    def test_choose_params_and_methods_raises(self, argnum):
+        """Test that the _choose_params_and_methods helper method raises an
+        error if incorrect trainable parameters are specified."""
+        tape = JacobianTape()
+        tape.trainable_params = [0]
+        diff_methods = ["F"]
+        with pytest.raises(
+            ValueError,
+            match="Incorrect trainable parameters",
+        ):
+            res = tape._choose_params_with_methods(diff_methods, argnum)
+
+    def test_choose_params_and_methods_warns_no_params(self):
+        """Test that the _choose_params_and_methods helper method warns if an
+        empty list was passed as argnum."""
+        tape = JacobianTape()
+        tape.trainable_params = [0]
+        diff_methods = ["F"]
+        argnum = []
+        with pytest.warns(
+            UserWarning,
+            match="No trainable parameters",
+        ):
+            res = tape._choose_params_with_methods(diff_methods, argnum)
+
 
 class TestJacobianIntegration:
     """Integration tests for the Jacobian method"""
@@ -471,6 +522,53 @@ class TestJacobianIntegration:
         expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]])
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    def test_single_expectation_value_with_argnum_all(self, tol):
+        """Tests correct output shape and evaluation for a tape
+        with a single expval output where all parameters are chose to compute
+        the jacobian"""
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        res = tape.jacobian(dev, argnum=[0, 1])  # <--- we choose both trainable parameters
+        assert res.shape == (1, 2)
+
+        expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_single_expectation_value_with_argnum_one(self, tol):
+        """Tests correct output shape and evaluation for a tape
+        with a single expval output where only one parameter is chosen to
+        estimate the jacobian.
+
+        This test relies on the fact that exactly one term of the estimated
+        jacobian will match the expected analytical value.
+        """
+        dev = qml.device("default.qubit", wires=2)
+        x = 0.543
+        y = -0.654
+
+        with JacobianTape() as tape:
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        res = tape.jacobian(dev, argnum=1)  # <--- we only choose one trainable parameter
+        assert res.shape == (1, 2)
+
+        expected = np.array([[0, np.cos(y) * np.cos(x)]])
+        res = res.flatten()
+        expected = expected.flatten()
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
     def test_multiple_expectation_values(self, tol):
         """Tests correct output shape and evaluation for a tape
         with multiple expval outputs"""
diff --git a/tests/tape/test_qnode.py b/tests/tape/test_qnode.py
index cacdedc25..aa3abe6fe 100644
--- a/tests/tape/test_qnode.py
+++ b/tests/tape/test_qnode.py
@@ -790,6 +790,40 @@ class TestIntegration:
 
         assert dev.num_executions == 6
 
+    @pytest.mark.parametrize("diff_method", ["parameter-shift", "finite-diff", "reversible"])
+    def test_single_expectation_value_with_argnum_one(self, diff_method, tol):
+        """Tests correct output shape and evaluation for a QNode
+        with a single expval output where only one parameter is chosen to
+        estimate the jacobian.
+
+        This test relies on the fact that exactly one term of the estimated
+        jacobian will match the expected analytical value.
+        """
+        from pennylane import numpy as anp
+
+        dev = qml.device("default.qubit", wires=2)
+
+        x = anp.array(0.543, requires_grad=True)
+        y = anp.array(-0.654, requires_grad=True)
+
+        @qml.qnode(
+            dev, diff_method=diff_method, argnum=[1]
+        )  # <--- we only choose one trainable parameter
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        res = qml.grad(circuit)(x, y)
+        assert len(res) == 2
+
+        expected = (0, np.cos(y) * np.cos(x))
+        res = res
+        expected = expected
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
 
 class TestMutability:
     """Test for QNode immutability"""