Allowing more flexible cost functions for optimizers

.github/CHANGELOG.md

@@ -2,6 +2,45 @@
 
 <h3>New features since last release</h3>
 
+* The built-in PennyLane optimizers allow more flexible cost functions. The cost function passed to most optimizers 
+  may accept any combination of trainable arguments, non-trainable arguments, and keyword arguments.
+  [(#959)](https://github.com/PennyLaneAI/pennylane/pull/959)
+
+  The full changes apply to:
+
+  * `AdagradOptimizer`
+  * `AdamOptimizer`
+  * `GradientDescentOptimizer`
+  * `MomentumOptimizer`
+  * `NesterovMomentumOptimizer`
+  * `RMSPropOptimizer`
+  * `RotosolveOptimizer` 
+
+  The `requires_grad=False` property must mark any non-trainable constant argument. 
+  The `RotoselectOptimizer` allows passing only keyword arguments.
+
+  Example use:
+
+  ```python
+  def cost(x, y, data, scale=1.0):
+      return scale * (x[0]-data)**2 + scale * (y-data)**2
+
+  x = np.array([1.], requires_grad=True)
+  y = np.array([1.0])
+  data = np.array([2.], requires_grad=False)
+
+  opt = qml.GradientDescentOptimizer()
+
+  # the optimizer step and step_and_cost methods can
+  # now update multiple parameters at once
+  x_new, y_new, data = opt.step(cost, x, y, data, scale=0.5)
+  (x_new, y_new, data), value = opt.step_and_cost(cost, x, y, data, scale=0.5) 
+
+  # list and tuple unpacking is also supported
+  params = (x, y, data)
+  params = opt.step(cost, *params)
+  ```
+
 * Support added for calculating the Hessian of quantum tapes using the second-order
   parameter shift formula.
   [(#961)](https://github.com/PennyLaneAI/pennylane/pull/961)
@@ -113,6 +152,7 @@
       return qml.expval(qml.PauliZ(0))
   ```
 
+
 <h3>Improvements</h3>
 
 * A new test series, pennylane/devices/tests/test_compare_default_qubit.py, has been added, allowing to test if
@@ -162,7 +202,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Olivia Di Matteo, Josh Izaac, Alejandro Montanez, Steven Oud, Chase Roberts, Maria Schuld, David Wierichs, Jiahao Yao.
+Olivia Di Matteo, Josh Izaac, Christina Lee, Alejandro Montanez, Steven Oud, Chase Roberts, Maria Schuld, David Wierichs, Jiahao Yao.
 
 # Release 0.13.0 (current release)
 

pennylane/optimize/adagrad.py

@@ -15,6 +15,7 @@
 import math
 
 from pennylane.utils import _flatten, unflatten
+from pennylane.numpy import ndarray, tensor
 from .gradient_descent import GradientDescentOptimizer
 
 
@@ -51,33 +52,62 @@ def __init__(self, stepsize=0.01, eps=1e-8):
         self.eps = eps
         self.accumulation = None
 
-    def apply_grad(self, grad, x):
-        r"""Update the variables x to take a single optimization step. Flattens and unflattens
+    def apply_grad(self, grad, args):
+        r"""Update the variables in args to take a single optimization step. Flattens and unflattens
         the inputs to maintain nested iterables as the parameters of the optimization.
 
         Args:
-            grad (array): The gradient of the objective
+            grad (tuple[array]): the gradient of the objective
                 function at point :math:`x^{(t)}`: :math:`\nabla f(x^{(t)})`
-            x (array): the current value of the variables :math:`x^{(t)}`
+            args (tuple): the current value of the variables :math:`x^{(t)}`
 
         Returns:
-            array: the new values :math:`x^{(t+1)}`
+            list: the new values :math:`x^{(t+1)}`
         """
-
-        x_flat = _flatten(x)
-        grad_flat = list(_flatten(grad))
+        args_new = list(args)
 
         if self.accumulation is None:
-            self.accumulation = [g * g for g in grad_flat]
-        else:
-            self.accumulation = [a + g * g for a, g in zip(self.accumulation, grad_flat)]
+            self.accumulation = [None] * len(args)
+
+        trained_index = 0
+        for index, arg in enumerate(args):
+            if getattr(arg, "requires_grad", True):
+                x_flat = _flatten(arg)
+                grad_flat = list(_flatten(grad[trained_index]))
+
+                self._update_accumulation(index, grad_flat)
+
+                x_new_flat = [
+                    e - (self._stepsize / math.sqrt(a + self.eps)) * g
+                    for a, g, e in zip(self.accumulation[index], grad_flat, x_flat)
+                ]
+
+                args_new[index] = unflatten(x_new_flat, arg)
 
-        x_new_flat = [
-            e - (self._stepsize / math.sqrt(a + self.eps)) * g
-            for a, g, e in zip(self.accumulation, grad_flat, x_flat)
-        ]
+                if isinstance(arg, ndarray):
+                    # Due to a bug in unflatten, input PennyLane tensors
+                    # are being unwrapped. Here, we cast them back to PennyLane
+                    # tensors.
+                    # TODO: remove when the following is fixed:
+                    # https://github.com/PennyLaneAI/pennylane/issues/966
+                    args_new[index] = args_new[index].view(tensor)
+                    args_new[index].requires_grad = True
 
-        return unflatten(x_new_flat, x)
+        return args_new
+
+    def _update_accumulation(self, index, grad_flat):
+        r"""Update the accumulation at index with gradient.
+
+        Args:
+            index (int): index of parameter to update.
+            grad_flat (list): flattened list form of gradient
+        """
+        if self.accumulation[index] is None:
+            self.accumulation[index] = [g * g for g in grad_flat]
+        else:
+            self.accumulation[index] = [
+                a + g * g for a, g in zip(self.accumulation[index], grad_flat)
+            ]
 
     def reset(self):
         """Reset optimizer by erasing memory of past steps."""

pennylane/optimize/adam.py

@@ -15,6 +15,7 @@
 import math
 
 from pennylane.utils import _flatten, unflatten
+from pennylane.numpy import ndarray, tensor
 from .gradient_descent import GradientDescentOptimizer
 
 
@@ -61,49 +62,80 @@ def __init__(self, stepsize=0.01, beta1=0.9, beta2=0.99, eps=1e-8):
         self.sm = None
         self.t = 0
 
-    def apply_grad(self, grad, x):
-        r"""Update the variables x to take a single optimization step. Flattens and unflattens
+    def apply_grad(self, grad, args):
+        r"""Update the variables args to take a single optimization step. Flattens and unflattens
         the inputs to maintain nested iterables as the parameters of the optimization.
 
         Args:
-            grad (array): The gradient of the objective
+            grad (tuple[array]): the gradient of the objective
                 function at point :math:`x^{(t)}`: :math:`\nabla f(x^{(t)})`
-            x (array): the current value of the variables :math:`x^{(t)}`
+            args (tuple): the current value of the variables :math:`x^{(t)}`
 
         Returns:
-            array: the new values :math:`x^{(t+1)}`
+            list: the new values :math:`x^{(t+1)}`
         """
-
+        args_new = list(args)
         self.t += 1
 
-        grad_flat = list(_flatten(grad))
-        x_flat = _flatten(x)
+        # Update step size (instead of correcting for bias)
+        new_stepsize = (
+            self._stepsize * math.sqrt(1 - self.beta2 ** self.t) / (1 - self.beta1 ** self.t)
+        )
 
-        # Update first moment
         if self.fm is None:
-            self.fm = grad_flat
-        else:
-            self.fm = [self.beta1 * f + (1 - self.beta1) * g for f, g in zip(self.fm, grad_flat)]
+            self.fm = [None] * len(args)
 
-        # Update second moment
         if self.sm is None:
-            self.sm = [g * g for g in grad_flat]
-        else:
-            self.sm = [
-                self.beta2 * f + (1 - self.beta2) * g * g for f, g in zip(self.sm, grad_flat)
-            ]
+            self.sm = [None] * len(args)
 
-        # Update step size (instead of correcting for bias)
-        new_stepsize = (
-            self._stepsize * math.sqrt(1 - self.beta2 ** self.t) / (1 - self.beta1 ** self.t)
-        )
+        trained_index = 0
+        for index, arg in enumerate(args):
+            if getattr(arg, "requires_grad", True):
+                x_flat = _flatten(arg)
+                grad_flat = list(_flatten(grad[trained_index]))
+                trained_index += 1
 
-        x_new_flat = [
-            e - new_stepsize * f / (math.sqrt(s) + self.eps)
-            for f, s, e in zip(self.fm, self.sm, x_flat)
-        ]
+                self._update_moments(index, grad_flat)
 
-        return unflatten(x_new_flat, x)
+                x_new_flat = [
+                    e - new_stepsize * f / (math.sqrt(s) + self.eps)
+                    for f, s, e in zip(self.fm[index], self.sm[index], x_flat)
+                ]
+                args_new[index] = unflatten(x_new_flat, arg)
+
+                if isinstance(arg, ndarray):
+                    # Due to a bug in unflatten, input PennyLane tensors
+                    # are being unwrapped. Here, we cast them back to PennyLane
+                    # tensors.
+                    # TODO: remove when the following is fixed:
+                    # https://github.com/PennyLaneAI/pennylane/issues/966
+                    args_new[index] = args_new[index].view(tensor)
+                    args_new[index].requires_grad = True
+
+        return args_new
+
+    def _update_moments(self, index, grad_flat):
+        r"""Update the moments.
+
+        Args:
+            index (int): the index of the argument to update
+            grad_flat (list): the flattened gradient for that trainable param
+        """
+        # update first moment
+        if self.fm[index] is None:
+            self.fm[index] = grad_flat
+        else:
+            self.fm[index] = [
+                self.beta1 * f + (1 - self.beta1) * g for f, g in zip(self.fm[index], grad_flat)
+            ]
+
+        # update second moment
+        if self.sm[index] is None:
+            self.sm[index] = [g * g for g in grad_flat]
+        else:
+            self.sm[index] = [
+                self.beta2 * f + (1 - self.beta2) * g * g for f, g in zip(self.sm[index], grad_flat)
+            ]
 
     def reset(self):
         """Reset optimizer by erasing memory of past steps."""