Repository: PennyLaneAI/pennylane

Issue:
Top-level Jacobian only works with integer argnum (not list)
This seems to be due to a bug in how we are interfacing with autograd, or -- less likely -- a bug in autograd itself.

If the user provides a quantum function `func` with multiple input arguments in its signature, we can only successfully call `qml.jacobian(func, argnum=0)`, `qml.jacobian(func, argnum=1)`, etc., but not `qml.jacobian(func, argnum=[0,1])` (raises some error inside autograd). 

Since we've coded `jacobian` to behave nicely if all the arguments are combined into an array, a suitable alternate usage is available. Would be nice to figure this bug out at some stage


Discussion:
As far as I can tell, this issue occurs directly in autograd's Jacobian function, and is a limitation we might not be able to overcome. For instance, consider the following minimal non-working example:
```python
from autograd import jacobian
from autograd import numpy as np

def f(x, y):
    return np.sin(x)**2+y, y**2

jf = jacobian(f, argnum=[0, 1])
jf(np.array([0.5, 0.1]), np.array([0.4, -0.4]))
```
results in the following exception:
```console
/home/josh/.local/lib/python3.5/site-packages/autograd/tracer.py:14: UserWarning: Output seems independent of input.
  warnings.warn("Output seems independent of input.")
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/.local/lib/python3.5/site-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
     56     try:
---> 57         return getattr(obj, method)(*args, **kwds)
     58 

TypeError: 'ArrayVSpace' object cannot be interpreted as an integer
```
However, modifying the above minimal example such that
```python
jf = jacobian(f, argnum=0)
jf = jacobian(f, argnum=1)
```
works as expected, both returning a 2x2x2 Jacobian.

My feeling is that `autograd.jacobian` was designed assuming the user will _always_ want the Jacobian of a single NumPy array input and output.
I sent your MWE to the one autograd developer I know. Maybe he can provide advice

Resolving patch:
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -238,8 +238,8 @@ def jacobian(func, argnum):
             a combination of quantum and classical nodes. The output of the computation
             must consist of a single NumPy array (if classical) or a tuple of
             expectation values (if a quantum node)
-        argnum (int): which argument to take the gradient
-            with respect to. If the argument is a NumPy array, then the Jacobian
+        argnum (int or Sequence[int]): which argument to take the gradient
+            with respect to. If a sequence is given, the Jacobian matrix
             corresponding to all input elements and all output elements is returned.
 
     Returns:
@@ -247,7 +247,10 @@ def jacobian(func, argnum):
         function with respect to the arguments in argnum
     """
     # pylint: disable=no-value-for-parameter
-    return _jacobian(func, argnum)
+    if isinstance(argnum, int):
+        return _jacobian(func, argnum)
+    return lambda *args, **kwargs: numpy.stack([_jacobian(func, arg)(*args, **kwargs) for arg in argnum]).T
+
 
 
 def version():


Test patch:
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -923,6 +923,10 @@ def circuit(x, y, z):
 
         assert np.allclose(self.expected_jacobian(a, b, c), res, atol=tol, rtol=0)
 
+        #compare with what we get if argnum is a list
+        res2 = qml.jacobian(circuit, argnum=[0, 1, 2])(a, b, c)
+        assert np.allclose(res, res2, atol=tol, rtol=0)
+
     def test_multiple_expectation_jacobian_array(self, tol):
         """Tests that qnodes using an array argument return correct gradients
         for multiple expectation values."""