Repository: PennyLaneAI/pennylane
Pull request #1081 (2021-02-11 06:20:07)

Fix qml.grad so that the returned gradient always matches the cost function return type if only a single argument is differentiated
**Note: this PR was previously approved as #1067, but was reverted until later in the release cycle. Do not merge now.**

**Context:**

I realised during the recent CERN demo that `autograd.grad` behaves differently depending if the `argnum` argument is a sequence or an integer:

```python
>>> import autograd
>>> from autograd import numpy as np
>>> autograd.grad(np.sin, argnum=0)(0.6)
0.8253356149096783
>>> autograd.grad(np.sin, argnum=[0])(0.6)
(array(0.82533561),)
```

We never take this into account in PennyLane, which can lead to gradients often being returned as tuple of size 1.

**Description of the Change:**

* If the inferred `argnum` argument is a sequence of length 1, only the contained integer is passed to `autograd.grad`.

**Benefits:**

* The return of `qml.grad` should now _always_ match the return type of the function being differentiated.

* A lot of superfluous `[0]` could be removed from the tests (e.g., `qml.grad(circuit)(x)[0]`)

**Possible Drawbacks:**

* The optimizers have to add in the redundant tuple where needed so that the gradient remains iterable. But this is a easy change, and leads to an intuitive UI, so maybe better?

* Some of the demos might need to be updated to remove unneeded `[0]` after gradient computations. But I think very few demos actually show the gradients.

**Related GitHub Issues:** n/a


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081?src=pr&el=h1) Report
> Merging [#1081](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081?src=pr&el=desc) (61d1a01) into [revert-qml-grad](https://codecov.io/gh/PennyLaneAI/pennylane/commit/f7fcb290d9b82d80d2fd7e23b4e02454515eb11f?el=desc) (f7fcb29) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081?src=pr&el=tree)

```diff
@@               Coverage Diff                @@
##           revert-qml-grad    #1081   +/-   ##
================================================
  Coverage            97.72%   97.72%           
================================================
  Files                  154      154           
  Lines                11652    11658    +6     
================================================
+ Hits                 11387    11393    +6     
  Misses                 265      265           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pennylane/\_grad.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081/diff?src=pr&el=tree#diff-cGVubnlsYW5lL19ncmFkLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/optimize/gradient\_descent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081/diff?src=pr&el=tree#diff-cGVubnlsYW5lL29wdGltaXplL2dyYWRpZW50X2Rlc2NlbnQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/optimize/nesterov\_momentum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081/diff?src=pr&el=tree#diff-cGVubnlsYW5lL29wdGltaXplL25lc3Rlcm92X21vbWVudHVtLnB5) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081?src=pr&el=footer). Last update [f7fcb29...61d1a01](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1081?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 128026bf8..547724a1b 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -118,6 +118,12 @@
   (5, 4, 4)
   ```
 
+* If only one argument to the function `qml.grad` has the `requires_grad` attribute
+  set to True, then the returned gradient will be a NumPy array, rather than a
+  tuple of length 1.
+  [(#1067)](https://github.com/PennyLaneAI/pennylane/pull/1067)
+  [(#1081)](https://github.com/PennyLaneAI/pennylane/pull/1081)
+
 * An improvement has been made to how `QubitDevice` generates and post-processess samples,
   allowing QNode measurement statistics to work on devices with more than 32 qubits.
   [(#1088)](https://github.com/PennyLaneAI/pennylane/pull/1088)
diff --git a/pennylane/_grad.py b/pennylane/_grad.py
index e0a8de87a..7f6a4b03f 100644
--- a/pennylane/_grad.py
+++ b/pennylane/_grad.py
@@ -85,6 +85,9 @@ class grad:
             if getattr(arg, "requires_grad", True):
                 argnum.append(idx)
 
+        if len(argnum) == 1:
+            argnum = argnum[0]
+
         return self._grad_with_forward(
             self._fun,
             argnum=argnum,
diff --git a/pennylane/optimize/gradient_descent.py b/pennylane/optimize/gradient_descent.py
index a175dbeba..17584f6f9 100644
--- a/pennylane/optimize/gradient_descent.py
+++ b/pennylane/optimize/gradient_descent.py
@@ -127,6 +127,9 @@ class GradientDescentOptimizer:
         grad = g(*args, **kwargs)
         forward = getattr(g, "forward", None)
 
+        if len(args) == 1:
+            grad = (grad,)
+
         return grad, forward
 
     def apply_grad(self, grad, args):
diff --git a/pennylane/optimize/nesterov_momentum.py b/pennylane/optimize/nesterov_momentum.py
index 08c1070ee..a561afee2 100644
--- a/pennylane/optimize/nesterov_momentum.py
+++ b/pennylane/optimize/nesterov_momentum.py
@@ -82,4 +82,7 @@ class NesterovMomentumOptimizer(MomentumOptimizer):
         grad = g(*shifted_args, **kwargs)
         forward = getattr(g, "forward", None)
 
+        if len(args) == 1:
+            grad = (grad,)
+
         return grad, forward


Test patch:
diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py
index 7f89b274e..c9b063d0b 100644
--- a/tests/devices/test_default_qubit_autograd.py
+++ b/tests/devices/test_default_qubit_autograd.py
@@ -352,7 +352,7 @@ class TestHighLevelIntegration:
             qml.init.strong_ent_layers_normal(n_wires=2, n_layers=2), requires_grad=True
         )
 
-        grad = qml.grad(circuit)(weights)[0]
+        grad = qml.grad(circuit)(weights)
         assert grad.shape == weights.shape
 
     def test_qnode_collection_integration(self):
@@ -374,7 +374,7 @@ class TestHighLevelIntegration:
         def cost(weights):
             return np.sum(qnodes(weights))
 
-        grad = qml.grad(cost)(weights)[0]
+        grad = qml.grad(cost)(weights)
         assert grad.shape == weights.shape
 
 class TestOps:
diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index 25111b313..782224e2e 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -541,8 +541,7 @@ class TestParameterHandlingIntegration:
 
         # we do not check for correctness, just that the output
         # is the correct shape
-        assert len(res) == 1
-        assert res[0].shape == weights.shape
+        assert res.shape == weights.shape
 
         # check that the first arg was marked as non-differentiable
         assert circuit.get_trainable_args() == {0}
@@ -587,8 +586,7 @@ class TestParameterHandlingIntegration:
 
         # we do not check for correctness, just that the output
         # is the correct shape
-        assert len(res) == 1
-        assert res[0].shape == weights.shape
+        assert res.shape == weights.shape
 
         # check that the second arg was marked as non-differentiable
         assert circuit.get_trainable_args() == {1}
@@ -633,8 +631,7 @@ class TestParameterHandlingIntegration:
 
         # we do not check for correctness, just that the output
         # is the correct shape
-        assert len(res) == 1
-        assert res[0].shape == weights.shape
+        assert res.shape == weights.shape
 
         # check that the last arg was marked as non-differentiable
         assert circuit.get_trainable_args() == {2}
@@ -748,7 +745,7 @@ class TestParameterHandlingIntegration:
         grad_fn = qml.grad(cost)
         res = grad_fn(weights)
 
-        assert len(res[0]) == 2
+        assert len(res) == 2
 
     def test_gradient_value(self, tol):
         """Test that the returned gradient value for a qubit QNode is correct,
diff --git a/tests/math/test_autograd_box.py b/tests/math/test_autograd_box.py
index f079a8086..249f5a92c 100644
--- a/tests/math/test_autograd_box.py
+++ b/tests/math/test_autograd_box.py
@@ -175,7 +175,7 @@ def test_autodifferentiation():
     cost_fn = lambda a: (qml.math.TensorBox(a).T() ** 2).unbox()[0, 1]
     grad_fn = qml.grad(cost_fn)
 
-    res = grad_fn(x)[0]
+    res = grad_fn(x)
     expected = np.array([[0.0, 0.0, 0.0], [8.0, 0.0, 0.0]])
     assert np.all(res == expected)
 
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 8480089b1..b3e55a791 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -975,7 +975,7 @@ class TestScatterElementAdd:
         assert isinstance(res, np.ndarray)
         assert fn.allclose(res, onp.array([[1.0, 1.0, 1.0], [1.0, 1.0, 1.3136]]))
 
-        grad = qml.grad(lambda weights: cost(weights)[1, 2])([x, y])[0]
+        grad = qml.grad(lambda weights: cost(weights)[1, 2])([x, y])
         assert fn.allclose(grad[0], onp.array([[0, 0, 0], [0, 0, 1.]]))
         assert fn.allclose(grad[1], 2 * y)
 
diff --git a/tests/tape/interfaces/test_qnode_autograd.py b/tests/tape/interfaces/test_qnode_autograd.py
index f7e70a15d..2e20b1105 100644
--- a/tests/tape/interfaces/test_qnode_autograd.py
+++ b/tests/tape/interfaces/test_qnode_autograd.py
@@ -118,9 +118,8 @@ class TestQNode:
 
         # gradients should work
         grad = qml.grad(circuit)(a)
-        assert len(grad) == 1
-        assert isinstance(grad[0], np.ndarray)
-        assert grad[0].shape == tuple()
+        assert isinstance(grad, float)
+        assert grad.shape == tuple()
 
     def test_interface_swap(self, dev_name, diff_method, tol):
         """Test that the autograd interface can be applied to a QNode
@@ -152,10 +151,9 @@ class TestQNode:
 
         res = circuit(a)
         grad = qml.grad(circuit)(a)
-        assert len(grad) == 1
 
         assert np.allclose(res, res_tf, atol=tol, rtol=0)
-        assert np.allclose(grad[0], grad_tf, atol=tol, rtol=0)
+        assert np.allclose(grad, grad_tf, atol=tol, rtol=0)
 
     def test_jacobian(self, dev_name, diff_method, mocker, tol):
         """Test jacobian calculation"""
@@ -594,6 +592,127 @@ class TestQNode:
         assert res.shape == (2, 10)
         assert isinstance(res, np.ndarray)
 
+    def test_gradient_non_differentiable_exception(self, dev_name, diff_method):
+        """Test that an exception is raised if non-differentiable data is
+        differentiated"""
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev, interface="autograd", diff_method=diff_method)
+        def circuit(data1):
+            qml.templates.AmplitudeEmbedding(data1, wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        grad_fn = qml.grad(circuit, argnum=0)
+        data1 = np.array([0, 1, 1, 0], requires_grad=False) / np.sqrt(2)
+
+        with pytest.raises(qml.numpy.NonDifferentiableError, match="is non-differentiable"):
+            grad_fn(data1)
+
+    def test_chained_qnodes(self, dev_name, diff_method):
+        """Test that the gradient of chained QNodes works without error"""
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev, interface="autograd", diff_method=diff_method)
+        def circuit1(weights):
+            qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
+
+        @qml.qnode(dev, interface="autograd", diff_method=diff_method)
+        def circuit2(data, weights):
+            qml.templates.AngleEmbedding(data, wires=[0, 1])
+            qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1])
+            return qml.expval(qml.PauliX(0))
+
+        def cost(weights):
+            w1, w2 = weights
+            c1 = circuit1(w1)
+            c2 = circuit2(c1, w2)
+            return np.sum(c2) ** 2
+
+        w1 = qml.init.strong_ent_layers_normal(n_wires=2, n_layers=3)
+        w2 = qml.init.strong_ent_layers_normal(n_wires=2, n_layers=4)
+
+        weights = [w1, w2]
+
+        grad_fn = qml.grad(cost)
+        res = grad_fn(weights)
+
+        assert len(res) == 2
+
+    def test_chained_gradient_value(self, dev_name, diff_method, tol):
+        """Test that the returned gradient value for two chained qubit QNodes
+        is correct."""
+        dev1 = qml.device(dev_name, wires=3)
+
+        @qml.qnode(dev1, diff_method=diff_method)
+        def circuit1(a, b, c):
+            qml.RX(a, wires=0)
+            qml.RX(b, wires=1)
+            qml.RX(c, wires=2)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(2))
+
+        dev2 = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev2, diff_method=diff_method)
+        def circuit2(data, weights):
+            qml.RX(data[0], wires=0)
+            qml.RX(data[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(weights[0], wires=0)
+            qml.RZ(weights[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(0) @ qml.PauliY(1))
+
+        def cost(a, b, c, weights):
+            return circuit2(circuit1(a, b, c), weights)
+
+        grad_fn = qml.grad(cost)
+
+        # Set the first parameter of circuit1 as non-differentiable.
+        a = np.array(0.4, requires_grad=False)
+
+        # The remaining free parameters are all differentiable.
+        b = 0.5
+        c = 0.1
+        weights = np.array([0.2, 0.3])
+
+        res = grad_fn(a, b, c, weights)
+
+        # Output should have shape [dcost/db, dcost/dc, dcost/dw],
+        # where b,c are scalars, and w is a vector of length 2.
+        assert len(res) == 3
+        assert res[0].shape == tuple() # scalar
+        assert res[1].shape == tuple() # scalar
+        assert res[2].shape == (2,)    # vector
+
+        cacbsc = np.cos(a)*np.cos(b)*np.sin(c)
+
+        expected = np.array([
+            # analytic expression for dcost/db
+            -np.cos(a)*np.sin(b)*np.sin(c)*np.cos(cacbsc)*np.sin(weights[0])*np.sin(np.cos(a)),
+            # analytic expression for dcost/dc
+            np.cos(a)*np.cos(b)*np.cos(c)*np.cos(cacbsc)*np.sin(weights[0])*np.sin(np.cos(a)),
+            # analytic expression for dcost/dw[0]
+            np.sin(cacbsc)*np.cos(weights[0])*np.sin(np.cos(a)),
+            # analytic expression for dcost/dw[1]
+            0
+        ])
+
+        # np.hstack 'flattens' the ragged gradient array allowing it
+        # to be compared with the expected result
+        assert np.allclose(np.hstack(res), expected, atol=tol, rtol=0)
+
+        if diff_method != "backprop":
+            # Check that the gradient was computed
+            # for all parameters in circuit2
+            assert circuit2.qtape.trainable_params == {0, 1, 2, 3}
+
+            # Check that the parameter-shift rule was not applied
+            # to the first parameter of circuit1.
+            assert circuit1.qtape.trainable_params == {1, 2}
+
 
 def qtransform(qnode, a, framework=np):
     """Transforms every RY(y) gate in a circuit to RX(-a*cos(y))"""
diff --git a/tests/test_classical_gradients.py b/tests/test_classical_gradients.py
index fc3aa44be..69dda25bc 100644
--- a/tests/test_classical_gradients.py
+++ b/tests/test_classical_gradients.py
@@ -306,7 +306,7 @@ class TestGrad:
         res = grad_fn(x, y)
         expected = np.array([np.cos(x) * np.cos(y) + y ** 2])
         assert np.allclose(res, expected, atol=tol, rtol=0)
-        assert spy.call_args_list[0][1]["argnum"] == [0]
+        assert spy.call_args_list[0][1]["argnum"] == 0
 
 
 class TestJacobian: