Fix how trainable args are counted for gradients in GradientDescentOptimizer and NesterovMomentumOptimizer

.github/CHANGELOG.md

@@ -462,6 +462,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)
+
 * Fixes a bug where the adjoint of `qml.QFT` when using the `qml.adjoint` function
   was not correctly computed.
   [(#1451)](https://github.com/PennyLaneAI/pennylane/pull/1451)

pennylane/optimize/gradient_descent.py

@@ -129,7 +129,12 @@ def compute_grad(objective_fn, args, kwargs, grad_fn=None):
         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

pennylane/optimize/nesterov_momentum.py

@@ -59,8 +59,13 @@ def compute_grad(self, objective_fn, args, kwargs, grad_fn=None):
         """
         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 @@ def compute_grad(self, objective_fn, args, kwargs, grad_fn=None):
         grad = g(*shifted_args, **kwargs)
         forward = getattr(g, "forward", None)
 
-        if len(args) == 1:
+        if len(trainable_args) == 1:
             grad = (grad,)
 
         return grad, forward

tests/test_optimize.py

@@ -877,3 +877,53 @@ def func2(args):
         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[0]) ** 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)
+        assert (
+            x == 0
+        )  # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
+        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[0]) ** 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)
+        assert (
+            x == 0
+        )  # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
+        assert ev == original_ev