Repository: PennyLaneAI/pennylane
Pull request #1699 (2021-09-29 08:58:20)

Fixes a bug where batch transforms wouldn't work with templates
**Context:** I discovered while hacking that batch transforms would not correctly apply to QNodes containing templates. This is because templates may perform classical computation, which the transform needs to take into account. However, the classical Jacobian computation was not using the same decomposition as the quantum component, resulting in a failure to perform the matrix multiplication due to a shape mismatch.

**Description of the Change:** Modifies the `classical_jacobian` function to use the same circuit decomposition as the quantum component.

**Benefits:** Batch transforms now work with templates.

**Possible Drawbacks:** n/a

**Related GitHub Issues:** n/a


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1699](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c4b549e) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/5d55cf8fe5e68fba42e6c52b42acde7fce3b2cf4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (5d55cf8) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

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

```diff
@@           Coverage Diff           @@
##           master    #1699   +/-   ##
=======================================
  Coverage   99.20%   99.20%           
=======================================
  Files         201      201           
  Lines       15122    15127    +5     
=======================================
+ Hits        15002    15007    +5     
  Misses        120      120           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/gradients/gradient\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWV0cmljX3RlbnNvci5weQ==) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699?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/1699?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [5d55cf8...c4b549e](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1699?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/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f9767d2ab..0584a6794 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -315,6 +315,10 @@
 * Fixes a bug in `default.mixed`, to ensure that returned probabilities are always non-negative.
   [(#1680)](https://github.com/PennyLaneAI/pennylane/pull/1680)
 
+* Fixes a bug where gradient transforms would fail to apply to QNodes
+  containing classical processing.
+  [(#1699)](https://github.com/PennyLaneAI/pennylane/pull/1699)
+
 <h3>Documentation</h3>
 
 * Adds a link to https://pennylane.ai/qml/demonstrations.html in the navbar.
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index b4a194fb8..169965979 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -150,7 +150,7 @@ class gradient_transform(qml.batch_transform):
         # inside the QNode.
         hybrid = tkwargs.pop("hybrid", self.hybrid)
         _wrapper = super().default_qnode_wrapper(qnode, targs, tkwargs)
-        cjac_fn = qml.transforms.classical_jacobian(qnode)
+        cjac_fn = qml.transforms.classical_jacobian(qnode, expand_fn=gradient_expand)
 
         def jacobian_wrapper(*args, **kwargs):
             qjac = _wrapper(*args, **kwargs)
diff --git a/pennylane/transforms/classical_jacobian.py b/pennylane/transforms/classical_jacobian.py
index 03b4b1154..44f30da46 100644
--- a/pennylane/transforms/classical_jacobian.py
+++ b/pennylane/transforms/classical_jacobian.py
@@ -19,7 +19,7 @@ import pennylane as qml
 from pennylane import numpy as np
 
 
-def classical_jacobian(qnode, argnum=None):
+def classical_jacobian(qnode, argnum=None, expand_fn=None):
     r"""Returns a function to extract the Jacobian
     matrix of the classical part of a QNode.
 
@@ -30,6 +30,8 @@ def classical_jacobian(qnode, argnum=None):
         qnode (pennylane.QNode): QNode to compute the (classical) Jacobian of
         argnum (int or Sequence[int]): indices of QNode arguments with respect to which
             the (classical) Jacobian is computed
+        expand_fn (None or function): an expansion function (if required) to be applied to the
+            QNode quantum tape before the classical Jacobian is computed
 
     Returns:
         function: Function which accepts the same arguments as the QNode.
@@ -137,7 +139,12 @@ def classical_jacobian(qnode, argnum=None):
         """Returns the trainable gate parameters for a given QNode input."""
         trainable_only = kwargs.pop("_trainable_only", True)
         qnode.construct(args, kwargs)
-        return qml.math.stack(qnode.qtape.get_parameters(trainable_only=trainable_only))
+        tape = qnode.qtape
+
+        if expand_fn is not None:
+            tape = expand_fn(tape)
+
+        return qml.math.stack(tape.get_parameters(trainable_only=trainable_only))
 
     if qnode.interface == "autograd":
 
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index aef150680..0a870d3cf 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -230,7 +230,11 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
         qnode = qnode.qnodes.qnodes[0]
 
     mt_fn = self.default_qnode_wrapper(qnode, targs, tkwargs)
-    cjac_fn = qml.transforms.classical_jacobian(qnode)
+
+    if isinstance(qnode, qml.beta.QNode):
+        cjac_fn = qml.transforms.classical_jacobian(qnode, expand_fn=self.expand_fn)
+    else:
+        cjac_fn = qml.transforms.classical_jacobian(qnode)
 
     def wrapper(*args, **kwargs):
         mt = mt_fn(*args, **kwargs)


Test patch:
diff --git a/tests/gradients/test_gradient_transform.py b/tests/gradients/test_gradient_transform.py
index 88670944e..a1882d4aa 100644
--- a/tests/gradients/test_gradient_transform.py
+++ b/tests/gradients/test_gradient_transform.py
@@ -329,6 +329,24 @@ class TestGradientTransformIntegration:
         expected = qml.jacobian(circuit)(w)
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("strategy", ["gradient", "device"])
+    def test_template_integration(self, strategy, tol):
+        """Test that the gradient transform acts on QNodes
+        correctly when the QNode contains a template"""
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.beta.qnode(dev, expansion_strategy=strategy)
+        def circuit(weights):
+            qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+            return qml.probs(wires=[0, 1])
+
+        weights = np.ones([2, 3, 3], dtype=np.float64, requires_grad=True)
+        res = qml.gradients.param_shift(circuit)(weights)
+        assert res.shape == (4, 2, 3, 3)
+
+        expected = qml.jacobian(circuit)(weights)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
 
 class TestInterfaceIntegration:
     """Test that the gradient transforms are differentiable
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index 76ff4370d..4bd6aa59d 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -243,6 +243,21 @@ class TestMetricTensor:
         )
         assert np.allclose(g, np.diag(expected), atol=tol, rtol=0)
 
+    @pytest.mark.parametrize("strategy", ["gradient", "device"])
+    def test_template_integration(self, strategy, tol):
+        """Test that the metric tensor transform acts on QNodes
+        correctly when the QNode contains a template"""
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.beta.qnode(dev, expansion_strategy=strategy)
+        def circuit(weights):
+            qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+            return qml.probs(wires=[0, 1])
+
+        weights = np.ones([2, 3, 3], dtype=np.float64, requires_grad=True)
+        res = qml.metric_tensor(circuit)(weights)
+        assert res.shape == (2, 3, 3, 2, 3, 3)
+
     def test_evaluate_diag_metric_tensor_classical_processing(self, tol):
         """Test that a diagonal metric tensor evaluates correctly
         when the QNode includes classical processing."""