Repository: PennyLaneAI/pennylane

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:
  ```
  



Discussion:
Not to be confused with the "Stochastic Parameter-shift Rule" of Banchi & Crooks? :thinking: 
Yes, not to be confused! I was thinking about that while writing this issue. Is there a better name for this ‘stochastic’ parameter shift rule?
maybe have a "stochastic gradient" mixin that applies to finite diff, PSR, etc? i.e., the stochasticity is independent of the actual method?
Great idea! It wouldn't even need to be a mixin, we can just add support directly to the `tape.jacobian()` method, which is diff method agnostic.
Hi @josh146,

Then, is it only adding the code that you provide above indicating that the number of k_params should be between 1 and the number of parameters? (like the code below)
```Python
        k_params = options.get("k_params", None)
        
        params_and_methods = enumerate(diff_methods)
        if k_params is not None:
            if not 1 <= k_params <= self.num_params:
                raise ValueError(f"The number of k_params: {k_params} is not between the range 1 and {self.num_params}")

            params_and_methods = set(random.choices(list(enumerate(diff_methods)), k=k_params))
            
        for trainable_idx, param_method in params_and_methods:
            if (method == "best" and param_method[0] == "F") or (method == "numeric"):
                # numeric method
                tapes, processing_fn = self.numeric_pd(trainable_idx, params=params, **options)

            elif (method == "best" and param_method[0] == "A") or (method == "analytic"):
                # analytic method
                tapes, processing_fn = self.analytic_pd(trainable_idx, params=params, **options)

            processing_fns.append(processing_fn)

            # we create a flat list here to feed at once to the device
            all_tapes.extend(tapes)

            # to extract the correct result for this parameter later, remember the number of tapes
            reshape_info.append(len(tapes))
```
I found a problem with this approach. It could select the following case for `param_method = '0'` which is the case for the param ` qml.RX(0.133, wires='a')`. Therefore, neither the case (method == "best" and param_method[0] == "F")  or (method == "best" and param_method[0] == "A") is selected.

```Python
import pennylane as qml
import pennylane.tape

qml.enable_tape()
dev = qml.device("default.qubit", wires=[0,'a'])

with qml.tape.JacobianTape() as tape:
    qml.RX(0.432, wires=0)
    qml.RY(0.543, wires=0)
    qml.RX(0.133, wires=0)
    qml.CNOT(wires=[0,'a'])
    qml.RX(0.133, wires='a')
    qml.expval(qml.PauliZ(wires=[0]))

opt = {'k_params':1}
tape.jacobian(dev, **opt)

---------------------------------------------------------------------------
UnboundLocalError                         Traceback (most recent call last)
<ipython-input-15-ae622e5dc974> in <module>
      1 opt = {'k_params':1}
      2 # tape.trainable_params = {3}
----> 3 tape.jacobian(dev, **opt)

~/pennylane/pennylane/tape/tapes/jacobian_tape.py in jacobian(self, device, params, **options)
    517                 tapes, processing_fn = self.analytic_pd(trainable_idx, params=params, **options)
    518 
--> 519             processing_fns.append(processing_fn)
    520 
    521             # we create a flat list here to feed at once to the device

UnboundLocalError: local variable 'processing_fn' referenced before assignment
```
Hi @alejomonbar! Would the following continue statement (which I believe is part of the existing logic) fix this?

```python
        for trainable_idx, param_method in params_and_methods:
            if param_method == "0":
                continue
```

Resolving patch:
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
--- 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
--- a/pennylane/tape/jacobian_tape.py
+++ b/pennylane/tape/jacobian_tape.py
@@ -18,10 +18,11 @@
 # 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 @@ def hessian_pd(self, i, j, params, **options):
         """
         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 @@ def jacobian(self, device, params=None, **options):
             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 @@ def jacobian(self, device, params=None, **options):
         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
--- a/tests/tape/test_jacobian_tape.py
+++ b/tests/tape/test_jacobian_tape.py
@@ -431,6 +431,57 @@ def test_independent_parameters(self):
         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 @@ def test_single_expectation_value(self, tol):
         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
--- a/tests/tape/test_qnode.py
+++ b/tests/tape/test_qnode.py
@@ -790,6 +790,40 @@ def func():
 
         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"""