Repository: PennyLaneAI/pennylane

Issue:
Device jacobian is not called when QNode contains gates that need to be decomposed
When a QNode contains gates that are not natively supported on a device they are decomposed by PennyLane (which is great!). 

Unfortunately (for some reason I do not understand) `diff_method='device'` does no longer work for QNodes containing such gates as is demonstrated by the following example:
```python
import pennylane as qml
from pennylane.devices import DefaultQubit
from pennylane.operation import Operation
import numpy as np

class MyDev(DefaultQubit):
    @classmethod
    def capabilities(cls):
        capabilities = super().capabilities().copy()
        capabilities.update(
            provides_jacobian=True,
        )
        return capabilities

    def jacobian(self, operations_or_tape, observables=None, parameters=None):
        raise Exception("This exception signals that the device jacobian() was called")

class RYZ(Operation):
    num_params = 2
    num_wires = 1
    par_domain = "R"
    grad_method = "A"

    @staticmethod
    def _matrix(*params):
        theta, phi = params
        return qml.RZ._matrix(phi + np.pi) @ qml.RY._matrix(theta + np.pi / 2)

    @staticmethod
    def decomposition(theta, phi, wires):
        decomp_ops = [
            qml.RY(np.pi / 2, wires=wires),
            qml.RY(theta, wires=wires),
            qml.RZ(phi, wires=wires),
            qml.RZ(np.pi, wires=wires),
        ]
        return decomp_ops

@qml.qnode(MyDev(wires=2), diff_method='device')
def circuit(params):
    RYZ(params[0], params[1], wires=0) # with this commented the device jacobian is called also with PennyLane master
    qml.RY(params[0], wires=0)
    return qml.expval(qml.Hermitian(np.array([[1, 0], [0, 0]]), wires=0))

params = np.array([0.75916255, -2.16921113])
qml.jacobian(circuit)(params)
raise Exception("This line should not be reached as device jacobian() should have been called")
```
With PL 0.13.0 `MyDev.jacobian()` is called, as is expected. With master `MyDev.jacobian()` is not called and PL computes the jacobian in some other way. By commenting out the `RYZ(params[0], params[1], wires=0)`, and thereby removing the only gate that needs to be decomposed, the device `jacobian()` is called again.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117?src=pr&el=h1) Report
> Merging [#1117](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117?src=pr&el=desc) (8f79097) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/1206dbda6c43807c2cdf08064c11bc9c82d75deb?el=desc) (1206dbd) will **not change** coverage.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##           master    #1117   +/-   ##
=======================================
  Coverage   97.74%   97.74%           
=======================================
  Files         154      154           
  Lines       11726    11726           
=======================================
  Hits        11461    11461           
  Misses        265      265           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pennylane/tape/qnode.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117/diff?src=pr&el=tree#diff-cGVubnlsYW5lL3RhcGUvcW5vZGUucHk=) | `98.93% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117?src=pr&el=footer). Last update [1206dbd...8f79097](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1117?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 8e87f7165..128026bf8 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -166,6 +166,10 @@
   with non-expectation measurement statistics.
   [(#1106)](https://github.com/PennyLaneAI/pennylane/pull/1106)
 
+* Fixes a bug where decompositions would reset the differentiation method
+  of a QNode.
+  [(#1117)](https://github.com/PennyLaneAI/pennylane/pull/1117)
+
 <h3>Documentation</h3>
 
 - Typos addressed in templates documentation.
diff --git a/pennylane/tape/qnode.py b/pennylane/tape/qnode.py
index 020b888bf..5a7826d1a 100644
--- a/pennylane/tape/qnode.py
+++ b/pennylane/tape/qnode.py
@@ -516,9 +516,6 @@ class QNode:
                         "Operator {} must act on all wires".format(obj.name)
                     )
 
-        # provide the jacobian options
-        self.qtape.jacobian_options = self.diff_options
-
         # pylint: disable=protected-access
         obs_on_same_wire = len(self.qtape._obs_sharing_wires) > 0
         ops_not_supported = any(
@@ -536,6 +533,9 @@ class QNode:
                 and self.device.supports_operation(obj.name),
             )
 
+        # provide the jacobian options
+        self.qtape.jacobian_options = self.diff_options
+
     def __call__(self, *args, **kwargs):
 
         # If shots specified in call but not in qfunc signature,


Test patch:
diff --git a/tests/tape/tapes/test_qnode.py b/tests/tape/tapes/test_qnode.py
index ce8de947b..686f5d2ba 100644
--- a/tests/tape/tapes/test_qnode.py
+++ b/tests/tape/tapes/test_qnode.py
@@ -307,6 +307,55 @@ class TestTapeConstruction:
 
         assert jac.shape == (4, 2)
 
+    def test_diff_method_expansion(self, monkeypatch, mocker):
+        """Test that a QNode with tape expansion during construction
+        preserves the differentiation method."""
+
+        class MyDev(qml.devices.DefaultQubit):
+            """Dummy device that supports device Jacobians"""
+
+            @classmethod
+            def capabilities(cls):
+                capabilities = super().capabilities().copy()
+                capabilities.update(
+                    provides_jacobian=True,
+                )
+                return capabilities
+
+            def jacobian(self, *args, **kwargs):
+                return np.zeros((2, 4))
+
+        dev = MyDev(wires=2)
+
+        def func(x, y):
+            # the U2 operation is not supported on default.qubit
+            # and is decomposed.
+            qml.U2(x, y, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=0)
+
+        qn = QNode(func, dev, diff_method="device", h=1e-8, order=2)
+
+        assert qn.diff_options["method"] == "device"
+        assert qn.diff_options["h"] == 1e-8
+        assert qn.diff_options["order"] == 2
+
+        x = 0.12
+        y = 0.54
+
+        spy = mocker.spy(JacobianTape, "expand")
+        res = qn(x, y)
+
+        spy.assert_called_once()
+        assert qn.qtape.jacobian_options["method"] == "device"
+        assert qn.qtape.jacobian_options["h"] == 1e-8
+        assert qn.qtape.jacobian_options["order"] == 2
+
+        spy = mocker.spy(JacobianTape, "jacobian")
+        jac = qml.jacobian(qn)(x, y)
+
+        assert spy.call_args_list[0][1]["method"] == "device"
+
     def test_returning_non_measurements(self):
         """Test that an exception is raised if a non-measurement
         is returned from the QNode."""