Repository: PennyLaneAI/pennylane
Pull request #1487 (2021-07-30 15:32:44)

added grad_fn kwarg to QNGOptimizer.step_and_cost and .step
The `step_and_cost` and `step` methods of `QNGOptimizer` currently do not accept the `grad_fn` argument like other optimizers.
In particular, this prevents a straight-forward use of JAX with QNG.

**Changes**
The kwarg `grad_fn` is accepted as explicit kwarg in the two above methods, passing them to `compute_grad`.

**Benefits**
Unifying the optimizer interfaces and enabling custom gradient function when using QNG.

**Possible drawbacks**
None (The order of existing kwargs was not changed but `grad_fn` was added as last kwarg, so that even when using the existing kwargs as positional args there is no breaking change.

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1487?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1487](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1487?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (ca822f1) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c9905a625c1edb5582fe2097b544239f0e169c87?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c9905a6) will **not change** coverage.
> The diff coverage is `100.00%`.

> :exclamation: Current head ca822f1 differs from pull request most recent head 84cfcc5. Consider uploading reports for the commit 84cfcc5 to get more accurate results
[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1487/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/1487?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@           Coverage Diff           @@
##           master    #1487   +/-   ##
=======================================
  Coverage   98.32%   98.32%           
=======================================
  Files         180      180           
  Lines       12701    12701           
=======================================
  Hits        12488    12488           
  Misses        213      213           
```


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

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1487?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/1487?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [c9905a6...84cfcc5](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1487?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 0a58eef3e..a10d3d199 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -380,6 +380,10 @@
 
 <h3>Improvements</h3>
 
+* The `step` and `step_and_cost` methods of `QNGOptimizer` now accept a custom `grad_fn`
+  keyword argument to use for gradient computations.
+  [(#1487)](https://github.com/PennyLaneAI/pennylane/pull/1487)
+
 * The precision used by `default.qubit.jax` now matches the float precision
   indicated by 
   ```python
diff --git a/pennylane/optimize/gradient_descent.py b/pennylane/optimize/gradient_descent.py
index ecd70a2a4..8d34f45ce 100644
--- a/pennylane/optimize/gradient_descent.py
+++ b/pennylane/optimize/gradient_descent.py
@@ -58,7 +58,8 @@ class GradientDescentOptimizer:
             grad_fn (function): optional gradient function of the
                 objective function with respect to the variables ``*args``.
                 If ``None``, the gradient function is computed automatically.
-                Must return the same shape of tuple [array] as the autograd derivative.
+                Must return a ``tuple[array]`` with the same number of elements as ``*args``.
+                Each array of the tuple should have the same shape as the corresponding argument.
             **kwargs : variable length of keyword arguments for the objective function
 
         Returns:
@@ -87,7 +88,8 @@ class GradientDescentOptimizer:
             grad_fn (function): optional gradient function of the
                 objective function with respect to the variables ``x``.
                 If ``None``, the gradient function is computed automatically.
-                Must return the same shape of tuple [array] as the autograd derivative.
+                Must return a ``tuple[array]`` with the same number of elements as ``*args``.
+                Each array of the tuple should have the same shape as the corresponding argument.
             **kwargs : variable length of keyword arguments for the objective function
 
         Returns:
diff --git a/pennylane/optimize/qng.py b/pennylane/optimize/qng.py
index b06f6fc0e..90ccd91a0 100644
--- a/pennylane/optimize/qng.py
+++ b/pennylane/optimize/qng.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Quantum natural gradient optimizer"""
 # pylint: disable=too-many-branches
+# pylint: disable=too-many-arguments
 
 import numpy as np
 
@@ -155,19 +156,27 @@ class QNGOptimizer(GradientDescentOptimizer):
         self.metric_tensor = None
         self.lam = lam
 
-    def step_and_cost(self, qnode, x, recompute_tensor=True, metric_tensor_fn=None):
+    def step_and_cost(
+        self, qnode, *args, grad_fn=None, recompute_tensor=True, metric_tensor_fn=None, **kwargs
+    ):
         """Update the parameter array :math:`x` with one step of the optimizer and return the
         corresponding objective function value prior to the step.
 
         Args:
             qnode (QNode): the QNode for optimization
-            x (array): NumPy array containing the current values of the variables to be updated
+            *args : variable length argument list for qnode
+            grad_fn (function): optional gradient function of the
+                qnode with respect to the variables ``*args``.
+                If ``None``, the gradient function is computed automatically.
+                Must return a ``tuple[array]`` with the same number of elements as ``*args``.
+                Each array of the tuple should have the same shape as the corresponding argument.
             recompute_tensor (bool): Whether or not the metric tensor should
                 be recomputed. If not, the metric tensor from the previous
                 optimization step is used.
             metric_tensor_fn (function): Optional metric tensor function
-                with respect to the variables ``x``.
+                with respect to the variables ``args``.
                 If ``None``, the metric tensor function is computed automatically.
+            **kwargs : variable length of keyword arguments for the qnode
 
         Returns:
             tuple: the new variable values :math:`x^{(t+1)}` and the objective function output
@@ -183,53 +192,70 @@ class QNGOptimizer(GradientDescentOptimizer):
 
         if recompute_tensor or self.metric_tensor is None:
             if metric_tensor_fn is None:
-                # pseudo-inverse metric tensor
-                self.metric_tensor = qml.metric_tensor(qnode, diag_approx=self.diag_approx)(x)
-            else:
-                self.metric_tensor = metric_tensor_fn(x)
+                metric_tensor_fn = qml.metric_tensor(qnode, diag_approx=self.diag_approx)
+
+            self.metric_tensor = metric_tensor_fn(*args, **kwargs)
             self.metric_tensor += self.lam * np.identity(self.metric_tensor.shape[0])
 
-        # The QNGOptimizer.step does not permit passing an external gradient function.
-        # Autograd will always calculate the gradient and `forward` will never be `None`.
-        g, forward = self.compute_grad(qnode, (x,), dict())
-        x_out = self.apply_grad(g, x)
-        return x_out, forward
+        g, forward = self.compute_grad(qnode, args, kwargs, grad_fn=grad_fn)
+        new_args = self.apply_grad(g, args)
+
+        if forward is None:
+            forward = qnode(*args, **kwargs)
+
+        # unwrap from list if one argument, cleaner return
+        if len(new_args) == 1:
+            return new_args[0], forward
+        return new_args, forward
 
     # pylint: disable=arguments-differ
-    def step(self, qnode, x, recompute_tensor=True, metric_tensor_fn=None):
+    def step(
+        self, qnode, *args, grad_fn=None, recompute_tensor=True, metric_tensor_fn=None, **kwargs
+    ):
         """Update the parameter array :math:`x` with one step of the optimizer.
 
         Args:
             qnode (QNode): the QNode for optimization
-            x (array): NumPy array containing the current values of the variables to be updated
+            *args : variable length argument list for qnode
+            grad_fn (function): optional gradient function of the
+                qnode with respect to the variables ``*args``.
+                If ``None``, the gradient function is computed automatically.
+                Must return a ``tuple[array]`` with the same number of elements as ``*args``.
+                Each array of the tuple should have the same shape as the corresponding argument.
             recompute_tensor (bool): Whether or not the metric tensor should
                 be recomputed. If not, the metric tensor from the previous
                 optimization step is used.
             metric_tensor_fn (function): Optional metric tensor function
-                with respect to the variables ``x``.
+                with respect to the variables ``args``.
                 If ``None``, the metric tensor function is computed automatically.
+            **kwargs : variable length of keyword arguments for the qnode
 
         Returns:
             array: the new variable values :math:`x^{(t+1)}`
         """
-        x_out, _ = self.step_and_cost(
-            qnode, x, recompute_tensor=recompute_tensor, metric_tensor_fn=metric_tensor_fn
+        new_args, _ = self.step_and_cost(
+            qnode,
+            *args,
+            grad_fn=grad_fn,
+            recompute_tensor=recompute_tensor,
+            metric_tensor_fn=metric_tensor_fn,
+            **kwargs,
         )
-        return x_out
+        return new_args
 
-    def apply_grad(self, grad, x):
+    def apply_grad(self, grad, args):
         r"""Update the parameter array :math:`x` for 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
                 function at point :math:`x^{(t)}`: :math:`\nabla f(x^{(t)})`
-            x (array): the current value of the variables :math:`x^{(t)}`
+            args (array): the current value of the variables :math:`x^{(t)}`
 
         Returns:
             array: the new values :math:`x^{(t+1)}`
         """
         grad_flat = np.array(list(_flatten(grad)))
-        x_flat = np.array(list(_flatten(x)))
+        x_flat = np.array(list(_flatten(args)))
         x_new_flat = x_flat - self._stepsize * np.linalg.solve(self.metric_tensor, grad_flat)
-        return unflatten(x_new_flat, x)
+        return unflatten(x_new_flat, args)


Test patch:
diff --git a/tests/test_optimize_qng.py b/tests/test_optimize_qng.py
index 3eb482e5a..2ab2fbe7c 100644
--- a/tests/test_optimize_qng.py
+++ b/tests/test_optimize_qng.py
@@ -50,8 +50,8 @@ class TestOptimize:
     """Test basic optimization integration"""
 
     def test_step_and_cost_autograd(self, tol):
-        """Test that the correct cost is returned via the step_and_cost method for the QNG
-        optimizer"""
+        """Test that the correct cost and step is returned via the
+        step_and_cost method for the QNG optimizer"""
         dev = qml.device("default.qubit", wires=1)
 
         @qml.qnode(dev)
@@ -63,10 +63,78 @@ class TestOptimize:
         var = np.array([0.011, 0.012])
         opt = qml.QNGOptimizer(stepsize=0.01)
 
-        _, res = opt.step_and_cost(circuit, var)
+        step1, res = opt.step_and_cost(circuit, var)
+        step2 = opt.step(circuit, var)
 
         expected = circuit(var)
+        expected_step = var - opt._stepsize * 4 * qml.grad(circuit)(var)
         assert np.all(res == expected)
+        assert np.allclose(step1, expected_step)
+        assert np.allclose(step2, expected_step)
+
+    def test_step_and_cost_with_grad_fn_grouped_input(self, tol):
+        """Test that the correct cost and update is returned via the step_and_cost
+        method for the QNG optimizer when providing an explicit grad_fn.
+        Using a circuit with a single input containing all parameters."""
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(params):
+            qml.RX(params[0], wires=0)
+            qml.RY(params[1], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        var = np.array([0.011, 0.012])
+        opt = qml.QNGOptimizer(stepsize=0.01)
+
+        # With autograd gradient function
+        grad_fn = qml.grad(circuit)
+        step1, cost1 = opt.step_and_cost(circuit, var, grad_fn=grad_fn)
+        step2 = opt.step(circuit, var, grad_fn=grad_fn)
+
+        # With more custom gradient function, forward has to be computed explicitly.
+        grad_fn = lambda param: np.array(qml.grad(circuit)(param))
+        step3, cost2 = opt.step_and_cost(circuit, var, grad_fn=grad_fn)
+        step4 = opt.step(circuit, var, grad_fn=grad_fn)
+        expected_step = var - opt._stepsize * 4 * grad_fn(var)
+        expected_cost = circuit(var)
+
+        for step in [step1, step2, step3, step3]:
+            assert np.allclose(step, expected_step)
+        assert np.isclose(cost1, expected_cost)
+        assert np.isclose(cost2, expected_cost)
+
+    def test_step_and_cost_with_grad_fn_split_input(self, tol):
+        """Test that the correct cost and update is returned via the step_and_cost
+        method for the QNG optimizer when providing an explicit grad_fn.
+        Using a circuit with multiple inputs containing the parameters."""
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(params_0, params_1):
+            qml.RX(params_0, wires=0)
+            qml.RY(params_1, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        var = np.array([0.011, 0.012])
+        opt = qml.QNGOptimizer(stepsize=0.01)
+
+        # With autograd gradient function
+        grad_fn = qml.grad(circuit)
+        step1, cost1 = opt.step_and_cost(circuit, *var, grad_fn=grad_fn)
+        step2 = opt.step(circuit, *var, grad_fn=grad_fn)
+
+        # With more custom gradient function, forward has to be computed explicitly.
+        grad_fn = lambda params_0, params_1: np.array(qml.grad(circuit)(params_0, params_1))
+        step3, cost2 = opt.step_and_cost(circuit, *var, grad_fn=grad_fn)
+        step4 = opt.step(circuit, *var, grad_fn=grad_fn)
+        expected_step = var - opt._stepsize * 4 * grad_fn(*var)
+        expected_cost = circuit(*var)
+
+        for step in [step1, step2, step3, step3]:
+            assert np.allclose(step, expected_step)
+        assert np.isclose(cost1, expected_cost)
+        assert np.isclose(cost2, expected_cost)
 
     def test_qubit_rotation(self, tol):
         """Test qubit rotation has the correct QNG value