Repository: PennyLaneAI/pennylane
Pull request #1495 (2021-08-04 15:35:05)

Fix how trainable args are counted for gradients in `GradientDescentOptimizer` and `NesterovMomentumOptimizer`
**Context**
The `GradientDescentOptimizer` and `NesterovMomentumOptimizer` optimizers seem to have an issue with a cost function that takes one trainable and one non-trainable argument:
```python
import pennylane as qml
from pennylane import numpy as np

dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev, diff_method='parameter-shift')
def circuit(x):
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

def cost(x, target):
    return circuit(x)

opt = qml.GradientDescentOptimizer(stepsize=10)
x = np.tensor(0.0, requires_grad=True)
ev = np.tensor(0.7781, requires_grad=False)

x, cost_val = opt.step_and_cost(cost, x, ev)
```
Raises
```
~/xanadu/pennylane/pennylane/optimize/gradient_descent.py in apply_grad(self, grad, args)
    153             if getattr(arg, "requires_grad", True):
    154                 x_flat = _flatten(arg)
--> 155                 grad_flat = _flatten(grad[trained_index])
    156                 trained_index += 1
    157 

IndexError: invalid index to scalar variable.
```
Although the cost function does not depend on `target`, the error is still raised. If we remove that argument, the gradient is computed well:

```python
import pennylane as qml
from pennylane import numpy as np

dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev, diff_method='parameter-shift')
def circuit(x):
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))


def cost(x):
    return circuit(x)

opt = qml.GradientDescentOptimizer(stepsize=10)
x = np.tensor(0.0, requires_grad=True)
ev = np.tensor(0.7781, requires_grad=False)

x, cost_val = opt.step_and_cost(cost, x)
```
For both optimizers, there is a part of the logic, where the output gradient is adjusted:
```
        if len(args) == 1:
            grad = (grad,)
```
**Changes made**

Changes how the output gradient is adjusted such that only the trainable arguments are considered.

**Related issues**

https://github.com/PennyLaneAI/qml/issues/309

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1495?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1495](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1495?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (d586c24) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/0428134c6918a2984157fbc2ed3af3d66fd3f696?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (0428134) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1495/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/1495?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1495   +/-   ##
=======================================
  Coverage   98.36%   98.36%           
=======================================
  Files         182      182           
  Lines       12933    12937    +4     
=======================================
+ Hits        12722    12726    +4     
  Misses        211      211           
```


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

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1495?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/1495?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [0428134...d586c24](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1495?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).


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 66ceb552f..d2e09ece1 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -466,6 +466,11 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug in `GradientDescentOptimizer` and `NesterovMomentumOptimizer`
+  where a cost function with one trainable parameter and non-trainable
+  parameters raised an error.
+  [(#1495)](https://github.com/PennyLaneAI/pennylane/pull/1495)
+
 * Fixed an example in the documentation's 
   [introduction to numpy gradients](https://pennylane.readthedocs.io/en/stable/introduction/interfaces/numpy.html), where 
   the wires were a non-differentiable argument to the QNode. 
diff --git a/pennylane/optimize/gradient_descent.py b/pennylane/optimize/gradient_descent.py
index 8d34f45ce..64be2881d 100644
--- a/pennylane/optimize/gradient_descent.py
+++ b/pennylane/optimize/gradient_descent.py
@@ -129,7 +129,12 @@ class GradientDescentOptimizer:
         grad = g(*args, **kwargs)
         forward = getattr(g, "forward", None)
 
-        if len(args) == 1:
+        num_trainable_args = 0
+        for arg in args:
+            if getattr(arg, "requires_grad", True):
+                num_trainable_args += 1
+
+        if num_trainable_args == 1:
             grad = (grad,)
 
         return grad, forward
diff --git a/pennylane/optimize/nesterov_momentum.py b/pennylane/optimize/nesterov_momentum.py
index f98af5792..884366037 100644
--- a/pennylane/optimize/nesterov_momentum.py
+++ b/pennylane/optimize/nesterov_momentum.py
@@ -59,8 +59,13 @@ class NesterovMomentumOptimizer(MomentumOptimizer):
         """
         shifted_args = list(args)
 
+        trainable_args = []
+        for arg in args:
+            if getattr(arg, "requires_grad", True):
+                trainable_args.append(arg)
+
         if self.accumulation:
-            for index, arg in enumerate(args):
+            for index, arg in enumerate(trainable_args):
                 if self.accumulation[index]:
                     x_flat = _flatten(arg)
                     acc = _flatten(self.accumulation[index])
@@ -82,7 +87,7 @@ class NesterovMomentumOptimizer(MomentumOptimizer):
         grad = g(*shifted_args, **kwargs)
         forward = getattr(g, "forward", None)
 
-        if len(args) == 1:
+        if len(trainable_args) == 1:
             grad = (grad,)
 
         return grad, forward


Test patch:
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index c96af732e..1c2b12ca1 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -747,16 +747,40 @@ def reset(opt):
         opt.reset()
 
 
+@pytest.fixture
+def opt(opt_name):
+    if opt_name == "gd":
+        return GradientDescentOptimizer(stepsize)
+
+    if opt_name == "nest":
+        return NesterovMomentumOptimizer(stepsize, momentum=gamma)
+
+    if opt_name == "moment":
+        return MomentumOptimizer(stepsize, momentum=gamma)
+
+    if opt_name == "ada":
+        return AdagradOptimizer(stepsize)
+
+    if opt_name == "rms":
+        return RMSPropOptimizer(stepsize, decay=gamma)
+
+    if opt_name == "adam":
+        return AdamOptimizer(stepsize, beta1=gamma, beta2=delta)
+
+    if opt_name == "roto":
+        return RotosolveOptimizer()
+
+
 @pytest.mark.parametrize(
-    "opt, opt_name",
+    "opt_name",
     [
-        (GradientDescentOptimizer(stepsize), "gd"),
-        (MomentumOptimizer(stepsize, momentum=gamma), "moment"),
-        (NesterovMomentumOptimizer(stepsize, momentum=gamma), "nest"),
-        (AdagradOptimizer(stepsize), "ada"),
-        (RMSPropOptimizer(stepsize, decay=gamma), "rms"),
-        (AdamOptimizer(stepsize, beta1=gamma, beta2=delta), "adam"),
-        (RotosolveOptimizer(), "roto"),
+        "gd",
+        "moment",
+        "nest",
+        "ada",
+        "rms",
+        "adam",
+        "roto",
     ],
 )
 class TestOverOpts:
@@ -877,3 +901,76 @@ class TestOverOpts:
         assert x_seperate == pytest.approx(args_new[0], abs=tol)
         assert y_seperate == pytest.approx(args_new[1], abs=tol)
         assert z_seperate == pytest.approx(args_new[2], abs=tol)
+
+    def test_one_trainable_one_non_trainable(self, opt, opt_name, tol):
+        """Tests that a cost function that takes one trainable and one
+        non-trainable parameter executes well."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        def cost(x, target):
+            return (circuit(x) - target) ** 2
+
+        ev = np.tensor(0.7781, requires_grad=False)
+        x = np.tensor(0.0, requires_grad=True)
+
+        original_ev = ev
+
+        (x, ev), cost = opt.step_and_cost(cost, x, ev)
+
+        # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
+        assert x == 0
+        assert ev == original_ev
+
+    def test_one_non_trainable_one_trainable(self, opt, opt_name, tol):
+        """Tests that a cost function that takes one non-trainable and one
+        trainable parameter executes well."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        def cost(target, x):  # Note: the order of the arguments has been swapped
+            return (circuit(x) - target) ** 2
+
+        ev = np.tensor(0.7781, requires_grad=False)
+        x = np.tensor(0.0, requires_grad=True)
+
+        original_ev = ev
+
+        (ev, x), cost = opt.step_and_cost(cost, ev, x)
+        # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
+        assert x == 0
+        assert ev == original_ev
+
+    def test_two_trainable_args(self, opt, opt_name, tol):
+        """Tests that a cost function that takes at least two trainable
+        arguments executes well."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=0)
+            qml.RX(y, wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        def cost(x, y, target):
+            return (circuit(x, y) - target) ** 2
+
+        ev = np.tensor(0.7781, requires_grad=False)
+        x = np.tensor(0.0, requires_grad=True)
+        y = np.tensor(0.0, requires_grad=True)
+
+        original_ev = ev
+
+        (x, y, ev), cost = opt.step_and_cost(cost, x, y, ev)
+
+        # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
+        assert x == 0
+        assert ev == original_ev