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:
Hi @cvjjm, this will likely be due to a change in the order of differentiation methods. As `MyDev` inherits from `DefaultQubit`, it prioritizes `backprop` over the device jacobian.

In PL `0.13.0`, the behaviour of master can be reproduced by using `qml.enable_tape()`.

For further details on the change, see the point in the `Improvements` section of the [Release notes of `v0.14.0`](https://pennylane.readthedocs.io/en/latest/development/release_notes.html#release-0-14-0) starting with
> The logic for choosing the ‘best’ differentiation method has been altered to improve performance.

Resolving patch:
diff --git a/pennylane/tape/qnode.py b/pennylane/tape/qnode.py
--- a/pennylane/tape/qnode.py
+++ b/pennylane/tape/qnode.py
@@ -516,9 +516,6 @@ def construct(self, args, kwargs):
                         "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 @@ def construct(self, args, kwargs):
                 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
--- a/tests/tape/tapes/test_qnode.py
+++ b/tests/tape/tapes/test_qnode.py
@@ -307,6 +307,55 @@ def func(x, y):
 
         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."""