Repository: PennyLaneAI/pennylane
Pull request #4890 (2023-11-28 15:36:04)

Integrate `device_vjp` with the QNode
This PR allows users to request device-provided VJP's at the QNode level.

```pycon
>>> dev = qml.device('default.qubit')
>>> @qml.qnode(dev, diff_method="adjoint", device_vjp=True)
>>> def circuit(x):
...     qml.RX(x, wires=0)
...     return qml.expval(qml.PauliZ(0))
>>> with dev.tracker:
...     g = qml.grad(circuit)(qml.numpy.array(0.1))
>>> dev.tracker.totals
{'batches': 1, 'simulations': 1, 'executions': 1, 'vjp_batches': 1, 'vjps': 1}
>>> g
-0.09983341664682815
```

[sc-51039]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4890?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`b77f674`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b77f67447541f2a816fe38e383536d3edb84edd7?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.65% compared to head [(`066d6fc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4890?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #4890      +/-   ##
==========================================
- Coverage   99.65%   99.64%   -0.02%     
==========================================
  Files         385      385              
  Lines       34896    34632     -264     
==========================================
- Hits        34776    34509     -267     
- Misses        120      123       +3     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4890?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: A quick question: will we be able to make using VJPs internally the default for the user (without needing to specify a keyword argument) _if the underlying diff method, device, and interface support it_?
albi3ro: > A quick question: will we be able to make using VJPs internally the default for the user (without needing to specify a keyword argument) _if the underlying diff method, device, and interface support it_?

Once we have good support for VJP's we can switch the default from `False` to `None`. `None` will indicate to pennylane "use it if you can".  But we have a few things to add before then.
albi3ro: > Thanks Christina! What is expected to change internally and externally after `device_vjp` is fully integrated? This looks good!

You can see the new workflow in the code for autograd and torch (`interfaces/autograd.py` and `interfaces/torch.py`). PR #4878  is moving over the code for jax (non-jit).

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 66347dda0..f2d034d36 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -160,12 +160,25 @@
 * `default.qubit` no longer uses a dense matrix for `MultiControlledX` for more than 8 operation wires.
   [(#4673)](https://github.com/PennyLaneAI/pennylane/pull/4673)
 
-* Autograd and torch can now use vjps provided by the device from the new device API. If a device provides
+* Autograd and torch can now use VJPs provided by the device from the new device API. If a device provides
   a vector Jacobian product, this can be selected by providing `device_vjp=True` to
-  `qml.execute`.
+  `qml.QNode` or `qml.execute`.
   [(#4557)](https://github.com/PennyLaneAI/pennylane/pull/4557)
   [(#4654)](https://github.com/PennyLaneAI/pennylane/pull/4654)
 
+```pycon
+>>> dev = qml.device('default.qubit')
+>>> @qml.qnode(dev, diff_method="adjoint", device_vjp=True)
+>>> def circuit(x):
+...     qml.RX(x, wires=0)
+...     return qml.expval(qml.PauliZ(0))
+>>> with dev.tracker:
+...     g = qml.grad(circuit)(qml.numpy.array(0.1))
+>>> dev.tracker.totals
+{'batches': 1, 'simulations': 1, 'executions': 1, 'vjp_batches': 1, 'vjps': 1}
+>>> g
+-0.09983341664682815
+```
 * Updates to some relevant Pytests to enable its use as a suite of benchmarks.
   [(#4703)](https://github.com/PennyLaneAI/pennylane/pull/4703)
 
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 0a4ac2584..920493781 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -172,6 +172,8 @@ class QNode:
             the maximum number of derivatives to support. Increasing this value allows
             for higher order derivatives to be extracted, at the cost of additional
             (classical) computational overhead during the backwards pass.
+        device_vjp (bool): Whether or not to use the device-provided Vector Jacobian Product (VJP).
+            A value of ``None`` indicates to use it if the device provides it, but use the full jacobian otherwise.
 
     Keyword Args:
         **kwargs: Any additional keyword arguments provided are passed to the differentiation
@@ -394,6 +396,7 @@ class QNode:
         cache=True,
         cachesize=10000,
         max_diff=1,
+        device_vjp=False,
         **gradient_kwargs,
     ):
         if logger.isEnabledFor(logging.DEBUG):
@@ -464,6 +467,7 @@ class QNode:
             "cachesize": cachesize,
             "max_diff": max_diff,
             "max_expansion": max_expansion,
+            "device_vjp": device_vjp,
         }
 
         if self.expansion_strategy == "device":
@@ -993,7 +997,7 @@ class QNode:
                 interface=self.interface,
                 gradient_method=_gradient_method,
                 grad_on_execution=None if grad_on_execution == "best" else grad_on_execution,
-                use_device_jacobian_product=False,
+                use_device_jacobian_product=self.execute_kwargs["device_vjp"],
             )
             device_transform_program, config = self.device.preprocess(execution_config=config)
             full_transform_program = self.transform_program + device_transform_program


Test patch:
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
index 33e26ea1c..b199cf02d 100644
--- a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
@@ -24,30 +24,32 @@ from pennylane import qnode
 from pennylane.devices import DefaultQubit
 
 qubit_device_and_diff_method = [
-    [DefaultQubit(), "finite-diff", False],
-    [DefaultQubit(), "parameter-shift", False],
-    [DefaultQubit(), "backprop", True],
-    [DefaultQubit(), "adjoint", True],
-    [DefaultQubit(), "adjoint", False],
-    [DefaultQubit(), "spsa", False],
-    [DefaultQubit(), "hadamard", False],
+    [qml.device("default.qubit"), "finite-diff", False, False],
+    [qml.device("default.qubit"), "parameter-shift", False, False],
+    [qml.device("default.qubit"), "backprop", True, False],
+    [qml.device("default.qubit"), "adjoint", True, False],
+    [qml.device("default.qubit"), "adjoint", False, False],
+    [qml.device("default.qubit"), "spsa", False, False],
+    [qml.device("default.qubit"), "hadamard", False, False],
 ]
 
 interface_qubit_device_and_diff_method = [
-    ["autograd", DefaultQubit(), "finite-diff", False],
-    ["autograd", DefaultQubit(), "parameter-shift", False],
-    ["autograd", DefaultQubit(), "backprop", True],
-    ["autograd", DefaultQubit(), "adjoint", True],
-    ["autograd", DefaultQubit(), "adjoint", False],
-    ["autograd", DefaultQubit(), "spsa", False],
-    ["autograd", DefaultQubit(), "hadamard", False],
-    ["auto", DefaultQubit(), "finite-diff", False],
-    ["auto", DefaultQubit(), "parameter-shift", False],
-    ["auto", DefaultQubit(), "backprop", True],
-    ["auto", DefaultQubit(), "adjoint", True],
-    ["auto", DefaultQubit(), "adjoint", False],
-    ["auto", DefaultQubit(), "spsa", False],
-    ["auto", DefaultQubit(), "hadamard", False],
+    ["autograd", DefaultQubit(), "finite-diff", False, False],
+    ["autograd", DefaultQubit(), "parameter-shift", False, False],
+    ["autograd", DefaultQubit(), "backprop", True, False],
+    ["autograd", DefaultQubit(), "adjoint", True, False],
+    ["autograd", DefaultQubit(), "adjoint", False, False],
+    ["autograd", DefaultQubit(), "adjoint", True, True],
+    ["autograd", DefaultQubit(), "adjoint", False, True],
+    ["autograd", DefaultQubit(), "spsa", False, False],
+    ["autograd", DefaultQubit(), "hadamard", False, False],
+    ["auto", DefaultQubit(), "finite-diff", False, False],
+    ["auto", DefaultQubit(), "parameter-shift", False, False],
+    ["auto", DefaultQubit(), "backprop", True, False],
+    ["auto", DefaultQubit(), "adjoint", True, False],
+    ["auto", DefaultQubit(), "adjoint", False, False],
+    ["auto", DefaultQubit(), "spsa", False, False],
+    ["auto", DefaultQubit(), "hadamard", False, False],
 ]
 
 pytestmark = pytest.mark.autograd
@@ -58,13 +60,16 @@ H_FOR_SPSA = 0.01
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_qubit_device_and_diff_method,
 )
 class TestQNode:
     """Test that using the QNode with Autograd integrates with the PennyLane stack"""
 
     # pylint:disable=unused-argument
-    def test_execution_no_interface(self, interface, dev, diff_method, grad_on_execution):
+    def test_execution_no_interface(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """Test execution works without an interface"""
         if diff_method == "backprop":
             pytest.skip("Test does not support backprop")
@@ -85,13 +90,19 @@ class TestQNode:
         assert isinstance(res, np.ndarray)
         assert res.shape == tuple()  # pylint: disable=comparison-with-callable
 
-    def test_execution_with_interface(self, interface, dev, diff_method, grad_on_execution):
+    def test_execution_with_interface(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """Test execution works with the interface"""
         if diff_method == "backprop":
             pytest.skip("Test does not support backprop")
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -106,10 +117,13 @@ class TestQNode:
         assert isinstance(grad, float)
         assert grad.shape == tuple()
 
-    def test_jacobian(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_jacobian(self, interface, dev, diff_method, grad_on_execution, tol, device_vjp):
         """Test jacobian calculation"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "spsa":
@@ -149,10 +163,15 @@ class TestQNode:
         assert res[1].shape == (2,)
         assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
-    def test_jacobian_no_evaluate(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_jacobian_no_evaluate(
+        self, interface, dev, diff_method, grad_on_execution, tol, device_vjp
+    ):
         """Test jacobian calculation when no prior circuit evaluation has been performed"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -186,14 +205,21 @@ class TestQNode:
         assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
         assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
-    def test_jacobian_options(self, interface, dev, diff_method, grad_on_execution):
+    def test_jacobian_options(self, interface, dev, diff_method, grad_on_execution, device_vjp):
         """Test setting jacobian options"""
         if diff_method == "backprop":
             pytest.skip("Test does not support backprop")
 
         a = np.array([0.1, 0.2], requires_grad=True)
 
-        @qnode(dev, interface=interface, h=1e-8, order=2, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface=interface,
+            h=1e-8,
+            order=2,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
+        )
         def circuit(a):
             qml.RY(a[0], wires=0)
             qml.RX(a[1], wires=0)
@@ -201,7 +227,9 @@ class TestQNode:
 
         qml.jacobian(circuit)(a)
 
-    def test_changing_trainability(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_changing_trainability(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test changing the trainability of parameters changes the
         number of differentiation requests made"""
         if diff_method != "parameter-shift":
@@ -215,6 +243,7 @@ class TestQNode:
             interface=interface,
             diff_method="parameter-shift",
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -252,14 +281,18 @@ class TestQNode:
         circuit(a, b)
         assert circuit.qtape.trainable_params == [1]
 
-    def test_classical_processing(self, interface, dev, diff_method, grad_on_execution):
+    def test_classical_processing(self, interface, dev, diff_method, grad_on_execution, device_vjp):
         """Test classical processing within the quantum tape"""
         a = np.array(0.1, requires_grad=True)
         b = np.array(0.2, requires_grad=False)
         c = np.array(0.3, requires_grad=True)
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b, c):
             qml.RY(a * c, wires=0)
@@ -277,11 +310,17 @@ class TestQNode:
         assert res[0].shape == ()
         assert res[1].shape == ()
 
-    def test_no_trainable_parameters(self, interface, dev, diff_method, grad_on_execution):
+    def test_no_trainable_parameters(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """Test evaluation and Jacobian if there are no trainable parameters"""
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -314,14 +353,20 @@ class TestQNode:
 
         assert grad == tuple()
 
-    def test_matrix_parameter(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_matrix_parameter(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test that the autograd interface works correctly
         with a matrix parameter"""
         U = np.array([[0, 1], [1, 0]], requires_grad=False)
         a = np.array(0.1, requires_grad=True)
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(U, a):
             qml.QubitUnitary(U, wires=0)
@@ -337,13 +382,17 @@ class TestQNode:
         assert np.allclose(res, np.sin(a), atol=tol, rtol=0)
 
     def test_gradient_non_differentiable_exception(
-        self, interface, dev, diff_method, grad_on_execution
+        self, interface, dev, diff_method, grad_on_execution, device_vjp
     ):
         """Test that an exception is raised if non-differentiable data is
         differentiated"""
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(data1):
             qml.templates.AmplitudeEmbedding(data1, wires=[0, 1])
@@ -355,11 +404,16 @@ class TestQNode:
         with pytest.raises(qml.numpy.NonDifferentiableError, match="is non-differentiable"):
             grad_fn(data1)
 
-    def test_differentiable_expand(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_differentiable_expand(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test that operation and nested tape expansion
         is differentiable"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "spsa":
@@ -488,16 +542,22 @@ class TestShotsIntegration:
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_qubit_device_and_diff_method,
 )
 class TestQubitIntegration:
     """Tests that ensure various qubit circuits integrate correctly"""
 
-    def test_probability_differentiation(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_probability_differentiation(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Tests correct output shape and evaluation for a tape
         with a single prob output"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "adjoint":
@@ -526,12 +586,15 @@ class TestQubitIntegration:
         assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected))
 
     def test_multiple_probability_differentiation(
-        self, interface, dev, diff_method, grad_on_execution, tol
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
     ):
         """Tests correct output shape and evaluation for a tape
         with multiple prob outputs"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "adjoint":
@@ -588,11 +651,16 @@ class TestQubitIntegration:
         )
         assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected))
 
-    def test_ragged_differentiation(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_ragged_differentiation(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "adjoint":
@@ -635,12 +703,15 @@ class TestQubitIntegration:
         assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_ragged_differentiation_variance(
-        self, interface, dev, diff_method, grad_on_execution, tol
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
     ):
         """Tests correct output shape and evaluation for a tape
         with prob and variance outputs"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "adjoint":
@@ -701,7 +772,7 @@ class TestQubitIntegration:
         # assert jac[1].shape == (3,)
         assert np.allclose(jac[1], expected[1], atol=tol, rtol=0)
 
-    def test_chained_qnodes(self, interface, dev, diff_method, grad_on_execution):
+    def test_chained_qnodes(self, interface, dev, diff_method, grad_on_execution, device_vjp):
         """Test that the gradient of chained QNodes works without error"""
 
         # pylint: disable=too-few-public-methods
@@ -719,7 +790,11 @@ class TestQubitIntegration:
             return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit2(data, weights):
             qml.templates.AngleEmbedding(data, wires=[0, 1])
@@ -744,11 +819,16 @@ class TestQubitIntegration:
 
         assert len(res) == 2
 
-    def test_chained_gradient_value(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_chained_gradient_value(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test that the returned gradient value for two chained qubit QNodes
         is correct."""
         kwargs = dict(
-            interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "spsa":
@@ -839,7 +919,9 @@ class TestQubitIntegration:
             # to the first parameter of circuit1.
             assert circuit1.qtape.trainable_params == [1, 2]
 
-    def test_second_derivative(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_second_derivative(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test second derivative calculation of a scalar valued QNode"""
         if diff_method not in {"parameter-shift", "backprop"}:
             pytest.skip("Test only supports parameter-shift or backprop")
@@ -850,6 +932,7 @@ class TestQubitIntegration:
             interface=interface,
             grad_on_execution=grad_on_execution,
             max_diff=2,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RY(x[0], wires=0)
@@ -879,7 +962,7 @@ class TestQubitIntegration:
 
         assert np.allclose(g2, expected_g2, atol=tol, rtol=0)
 
-    def test_hessian(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test hessian calculation of a scalar valued QNode"""
         if diff_method not in {"parameter-shift", "backprop"}:
             pytest.skip("Test only supports parameter-shift or backprop")
@@ -889,6 +972,7 @@ class TestQubitIntegration:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
         )
         def circuit(x):
@@ -931,7 +1015,9 @@ class TestQubitIntegration:
 
         assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
-    def test_hessian_unused_parameter(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_hessian_unused_parameter(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test hessian calculation of a scalar valued QNode"""
         if diff_method not in {"parameter-shift", "backprop"}:
             pytest.skip("Test only supports parameter-shift or backprop")
@@ -941,6 +1027,7 @@ class TestQubitIntegration:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
         )
         def circuit(x):
@@ -969,7 +1056,9 @@ class TestQubitIntegration:
 
         assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
-    def test_hessian_vector_valued(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_hessian_vector_valued(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test hessian calculation of a vector valued QNode"""
 
         if diff_method not in {"parameter-shift", "backprop"}:
@@ -980,6 +1069,7 @@ class TestQubitIntegration:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
         )
         def circuit(x):
@@ -1032,7 +1122,7 @@ class TestQubitIntegration:
         assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
     def test_hessian_vector_valued_postprocessing(
-        self, interface, dev, diff_method, grad_on_execution, tol
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
     ):
         """Test hessian calculation of a vector valued QNode with post-processing"""
         if diff_method not in {"parameter-shift", "backprop"}:
@@ -1043,6 +1133,7 @@ class TestQubitIntegration:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
         )
         def circuit(x):
@@ -1085,7 +1176,7 @@ class TestQubitIntegration:
         assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
     def test_hessian_vector_valued_separate_args(
-        self, interface, dev, diff_method, grad_on_execution, tol
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
     ):
         """Test hessian calculation of a vector valued QNode that has separate input arguments"""
         if diff_method not in {"parameter-shift", "backprop"}:
@@ -1096,6 +1187,7 @@ class TestQubitIntegration:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
         )
         def circuit(a, b):
@@ -1153,7 +1245,7 @@ class TestQubitIntegration:
             assert np.allclose(hess[0], exp_hess[0], atol=tol, rtol=0)
             assert np.allclose(hess[1], exp_hess[1], atol=tol, rtol=0)
 
-    def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test hessian calculation of a ragged QNode"""
         if diff_method not in {"parameter-shift", "backprop"}:
             pytest.skip("Test only supports parameter-shift or backprop")
@@ -1163,6 +1255,7 @@ class TestQubitIntegration:
             diff_method=diff_method,
             interface=interface,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
             max_diff=2,
         )
         def circuit(x):
@@ -1211,7 +1304,7 @@ class TestQubitIntegration:
 
         assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
 
-    def test_state(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test that the state can be returned and differentiated"""
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support states")
@@ -1220,7 +1313,11 @@ class TestQubitIntegration:
         y = np.array(-0.654, requires_grad=True)
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1251,10 +1348,15 @@ class TestQubitIntegration:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("state", [[1], [0, 1]])  # Basis state and state vector
-    def test_projector(self, state, interface, dev, diff_method, grad_on_execution, tol):
+    def test_projector(
+        self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test that the variance of a projector is correctly returned"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
 
         if diff_method == "adjoint":
@@ -1303,13 +1405,17 @@ class TestQubitIntegration:
         assert np.allclose(jac, expected, atol=tol, rtol=0)
 
 
-@pytest.mark.parametrize("dev,diff_method,grad_on_execution", qubit_device_and_diff_method)
+@pytest.mark.parametrize(
+    "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method
+)
 class TestTapeExpansion:
     """Test that tape expansion within the QNode integrates correctly
     with the Autograd interface"""
 
     @pytest.mark.parametrize("max_diff", [1, 2])
-    def test_gradient_expansion_trainable_only(self, dev, diff_method, grad_on_execution, max_diff):
+    def test_gradient_expansion_trainable_only(
+        self, dev, diff_method, grad_on_execution, max_diff, device_vjp
+    ):
         """Test that a *supported* operation with no gradient recipe is only
         expanded for parameter-shift and finite-differences when it is trainable."""
         if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"):
@@ -1326,7 +1432,13 @@ class TestTapeExpansion:
             def decomposition(self):
                 return [qml.RY(3 * self.data[0], wires=self.wires)]
 
-        @qnode(dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff)
+        @qnode(
+            dev,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            max_diff=max_diff,
+            device_vjp=device_vjp,
+        )
         def circuit(x, y):
             qml.Hadamard(wires=0)
             PhaseShift(x, wires=0)
@@ -1341,12 +1453,15 @@ class TestTapeExpansion:
 
     @pytest.mark.parametrize("max_diff", [1, 2])
     def test_hamiltonian_expansion_analytic(
-        self, dev, diff_method, grad_on_execution, max_diff, tol
+        self, dev, diff_method, grad_on_execution, max_diff, tol, device_vjp
     ):
         """Test that if there are non-commuting groups and the number of shots is None
         the first and second order gradients are correctly evaluated"""
         kwargs = dict(
-            diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            max_diff=max_diff,
+            device_vjp=device_vjp,
         )
 
         if diff_method in ["adjoint", "hadamard"]:
@@ -1403,7 +1518,9 @@ class TestTapeExpansion:
 
     @pytest.mark.slow
     @pytest.mark.parametrize("max_diff", [1, 2])
-    def test_hamiltonian_finite_shots(self, dev, diff_method, grad_on_execution, max_diff):
+    def test_hamiltonian_finite_shots(
+        self, dev, diff_method, grad_on_execution, max_diff, device_vjp
+    ):
         """Test that the Hamiltonian is correctly measured if there
         are non-commuting groups and the number of shots is finite
         and the first and second order gradients are correctly evaluated"""
@@ -1428,6 +1545,7 @@ class TestTapeExpansion:
             diff_method=diff_method,
             grad_on_execution=grad_on_execution,
             max_diff=max_diff,
+            device_vjp=device_vjp,
             **gradient_kwargs,
         )
         def circuit(data, weights, coeffs):
@@ -1579,15 +1697,21 @@ class TestSample:
         assert isinstance(result[2], np.ndarray)
 
 
-@pytest.mark.parametrize("dev,diff_method,grad_on_execution", qubit_device_and_diff_method)
+@pytest.mark.parametrize(
+    "dev,diff_method,grad_on_execution,device_vjp", qubit_device_and_diff_method
+)
 class TestReturn:
     """Class to test the shape of the Grad/Jacobian/Hessian with different return types."""
 
-    def test_grad_single_measurement_param(self, dev, diff_method, grad_on_execution):
+    def test_grad_single_measurement_param(self, dev, diff_method, grad_on_execution, device_vjp):
         """For one measurement and one param, the gradient is a float."""
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -1600,11 +1724,17 @@ class TestReturn:
 
         assert isinstance(grad, np.tensor if diff_method == "backprop" else float)
 
-    def test_grad_single_measurement_multiple_param(self, dev, diff_method, grad_on_execution):
+    def test_grad_single_measurement_multiple_param(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """For one measurement and multiple param, the gradient is a tuple of arrays."""
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -1622,12 +1752,16 @@ class TestReturn:
         assert grad[1].shape == ()
 
     def test_grad_single_measurement_multiple_param_array(
-        self, dev, diff_method, grad_on_execution
+        self, dev, diff_method, grad_on_execution, device_vjp
     ):
         """For one measurement and multiple param as a single array params, the gradient is an array."""
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -1642,14 +1776,20 @@ class TestReturn:
         assert len(grad) == 2
         assert grad.shape == (2,)
 
-    def test_jacobian_single_measurement_param_probs(self, dev, diff_method, grad_on_execution):
+    def test_jacobian_single_measurement_param_probs(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """For a multi dimensional measurement (probs), check that a single array is returned with the correct
         dimension"""
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -1664,7 +1804,7 @@ class TestReturn:
         assert jac.shape == (4,)
 
     def test_jacobian_single_measurement_probs_multiple_param(
-        self, dev, diff_method, grad_on_execution
+        self, dev, diff_method, grad_on_execution, device_vjp
     ):
         """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
         the correct dimension"""
@@ -1672,7 +1812,11 @@ class TestReturn:
             pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -1693,14 +1837,18 @@ class TestReturn:
         assert jac[1].shape == (4,)
 
     def test_jacobian_single_measurement_probs_multiple_param_single_array(
-        self, dev, diff_method, grad_on_execution
+        self, dev, diff_method, grad_on_execution, device_vjp
     ):
         """For a multi dimensional measurement (probs), check that a single array is returned."""
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -1713,14 +1861,20 @@ class TestReturn:
         assert isinstance(jac, np.ndarray)
         assert jac.shape == (4, 2)
 
-    def test_jacobian_multiple_measurement_single_param(self, dev, diff_method, grad_on_execution):
+    def test_jacobian_multiple_measurement_single_param(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """The jacobian of multiple measurements with a single params return an array."""
 
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -1738,7 +1892,7 @@ class TestReturn:
         assert jac.shape == (5,)
 
     def test_jacobian_multiple_measurement_multiple_param(
-        self, dev, diff_method, grad_on_execution
+        self, dev, diff_method, grad_on_execution, device_vjp
     ):
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
 
@@ -1746,7 +1900,11 @@ class TestReturn:
             pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -1771,7 +1929,7 @@ class TestReturn:
         assert jac[1].shape == (5,)
 
     def test_jacobian_multiple_measurement_multiple_param_array(
-        self, dev, diff_method, grad_on_execution
+        self, dev, diff_method, grad_on_execution, device_vjp
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
 
@@ -1779,7 +1937,11 @@ class TestReturn:
             pytest.skip("The adjoint method does not currently support returning probabilities")
 
         @qnode(
-            dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface="autograd",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -1796,7 +1958,7 @@ class TestReturn:
         assert isinstance(jac, np.ndarray)
         assert jac.shape == (5, 2)
 
-    def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_execution):
+    def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp):
         """The hessian of single a measurement with multiple params return a tuple of arrays."""
 
         if diff_method == "adjoint":
@@ -1811,6 +1973,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1832,7 +1995,9 @@ class TestReturn:
         assert isinstance(hess[1], np.ndarray)
         assert hess[1].shape == (2,)
 
-    def test_hessian_expval_multiple_param_array(self, dev, diff_method, grad_on_execution):
+    def test_hessian_expval_multiple_param_array(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """The hessian of single measurement with a multiple params array return a single array."""
 
         if diff_method == "adjoint":
@@ -1846,6 +2011,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -1858,7 +2024,7 @@ class TestReturn:
         assert isinstance(hess, np.ndarray)
         assert hess.shape == (2, 2)
 
-    def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution):
+    def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp):
         """The hessian of single a measurement with multiple params return a tuple of arrays."""
 
         if diff_method == "adjoint":
@@ -1875,6 +2041,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1896,7 +2063,9 @@ class TestReturn:
         assert isinstance(hess[1], np.ndarray)
         assert hess[1].shape == (2,)
 
-    def test_hessian_var_multiple_param_array(self, dev, diff_method, grad_on_execution):
+    def test_hessian_var_multiple_param_array(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """The hessian of single measurement with a multiple params array return a single array."""
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not currently support second-order diff.")
@@ -1911,6 +2080,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -1923,7 +2093,9 @@ class TestReturn:
         assert isinstance(hess, np.ndarray)
         assert hess.shape == (2, 2)
 
-    def test_hessian_probs_expval_multiple_params(self, dev, diff_method, grad_on_execution):
+    def test_hessian_probs_expval_multiple_params(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
         if diff_method in ["adjoint", "hadamard"]:
             pytest.skip("The adjoint method does not currently support second-order diff.")
@@ -1937,6 +2109,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1961,7 +2134,9 @@ class TestReturn:
         assert isinstance(hess[1], np.ndarray)
         assert hess[1].shape == (6,)
 
-    def test_hessian_expval_probs_multiple_param_array(self, dev, diff_method, grad_on_execution):
+    def test_hessian_expval_probs_multiple_param_array(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """The hessian of multiple measurements with a multiple param array return a single array."""
 
         if diff_method in ["adjoint", "hadamard"]:
@@ -1975,6 +2150,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -1990,7 +2166,9 @@ class TestReturn:
         assert isinstance(hess, np.ndarray)
         assert hess.shape == (3, 2, 2)  # pylint: disable=no-member
 
-    def test_hessian_probs_var_multiple_params(self, dev, diff_method, grad_on_execution):
+    def test_hessian_probs_var_multiple_params(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
 
         if diff_method == "adjoint":
@@ -2007,6 +2185,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -2031,7 +2210,9 @@ class TestReturn:
         assert isinstance(hess[1], np.ndarray)
         assert hess[1].shape == (6,)
 
-    def test_hessian_var_multiple_param_array2(self, dev, diff_method, grad_on_execution):
+    def test_hessian_var_multiple_param_array2(
+        self, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """The hessian of multiple measurements with a multiple param array return a single array."""
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not currently support second-order diff.")
@@ -2046,6 +2227,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
index 7a18a6e7b..c2020c2d0 100644
--- a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
@@ -27,13 +27,15 @@ jacobian = torch.autograd.functional.jacobian
 hessian = torch.autograd.functional.hessian
 
 qubit_device_and_diff_method = [
-    [DefaultQubit(), "finite-diff", False],
-    [DefaultQubit(), "parameter-shift", False],
-    [DefaultQubit(), "backprop", True],
-    [DefaultQubit(), "adjoint", True],
-    [DefaultQubit(), "adjoint", False],
-    [DefaultQubit(), "spsa", False],
-    [DefaultQubit(), "hadamard", False],
+    [DefaultQubit(), "finite-diff", False, False],
+    [DefaultQubit(), "parameter-shift", False, False],
+    [DefaultQubit(), "backprop", True, False],
+    [DefaultQubit(), "adjoint", True, False],
+    [DefaultQubit(), "adjoint", False, False],
+    [DefaultQubit(), "adjoint", True, True],
+    [DefaultQubit(), "adjoint", False, True],
+    [DefaultQubit(), "spsa", False, False],
+    [DefaultQubit(), "hadamard", False, False],
 ]
 
 interface_and_qubit_device_and_diff_method = [
@@ -46,18 +48,25 @@ H_FOR_SPSA = 0.01
 
 
 @pytest.mark.parametrize(
-    "interface, dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface, dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestQNode:
     """Test that using the QNode with Torch integrates with the PennyLane stack"""
 
-    def test_execution_with_interface(self, interface, dev, diff_method, grad_on_execution):
+    def test_execution_with_interface(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp
+    ):
         """Test execution works with the interface"""
         if diff_method == "backprop":
             pytest.skip("Test does not support backprop")
 
         @qnode(
-            dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
+            dev,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            interface=interface,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -84,12 +93,16 @@ class TestQNode:
         assert isinstance(grad, torch.Tensor)
         assert grad.shape == ()
 
-    def test_interface_swap(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_interface_swap(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test that the Torch interface can be applied to a QNode
         with a pre-existing interface"""
 
         @qnode(
-            dev, diff_method=diff_method, interface="autograd", grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface="autograd",
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a, wires=0)
@@ -115,7 +128,7 @@ class TestQNode:
         assert np.allclose(res1, res2.detach().numpy(), atol=tol, rtol=0)
         assert np.allclose(grad1, grad2, atol=tol, rtol=0)
 
-    def test_drawing(self, interface, dev, diff_method, grad_on_execution):
+    def test_drawing(self, interface, dev, diff_method, grad_on_execution, device_vjp):
         """Test circuit drawing when using the torch interface"""
 
         x = torch.tensor(0.1, requires_grad=True)
@@ -123,7 +136,11 @@ class TestQNode:
         z = torch.tensor(0.4, requires_grad=True)
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(p1, p2=y, **kwargs):
             qml.RX(p1, wires=0)
@@ -139,10 +156,13 @@ class TestQNode:
 
         assert result == expected
 
-    def test_jacobian(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_jacobian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test jacobian calculation"""
         kwargs = dict(
-            diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            interface=interface,
+            device_vjp=device_vjp,
         )
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -191,7 +211,14 @@ class TestQNode:
 
     # TODO: fix this behavior with float: already present before return type.
     @pytest.mark.xfail
-    def test_jacobian_dtype(self, interface, dev, diff_method, grad_on_execution):
+    def test_jacobian_dtype(
+        self,
+        interface,
+        dev,
+        diff_method,
+        grad_on_execution,
+        device_vjp,
+    ):
         """Test calculating the jacobian with a different datatype"""
         if diff_method == "backprop":
             pytest.skip("Test does not support backprop")
@@ -200,7 +227,11 @@ class TestQNode:
         b = torch.tensor(0.2, dtype=torch.float32, requires_grad=True)
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -224,7 +255,14 @@ class TestQNode:
         assert a.grad.dtype is torch.float32
         assert b.grad.dtype is torch.float32
 
-    def test_jacobian_options(self, interface, dev, diff_method, grad_on_execution):
+    def test_jacobian_options(
+        self,
+        interface,
+        dev,
+        diff_method,
+        grad_on_execution,
+        device_vjp,
+    ):
         """Test setting jacobian options"""
         if diff_method not in {"finite-diff", "spsa"}:
             pytest.skip("Test only works with finite-diff and spsa")
@@ -238,6 +276,7 @@ class TestQNode:
             interface=interface,
             h=1e-8,
             approx_order=2,
+            device_vjp=device_vjp,
         )
         def circuit(a):
             qml.RY(a[0], wires=0)
@@ -247,7 +286,9 @@ class TestQNode:
         res = circuit(a)
         res.backward()
 
-    def test_changing_trainability(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_changing_trainability(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test that changing the trainability of parameters changes the
         number of differentiation requests made"""
         if diff_method != "parameter-shift":
@@ -260,7 +301,11 @@ class TestQNode:
         b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True)
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -309,14 +354,25 @@ class TestQNode:
         expected = -np.sin(a_val) + np.sin(a_val) * np.sin(b_val)
         assert np.allclose(a.grad, expected, atol=tol, rtol=0)
 
-    def test_classical_processing(self, interface, dev, diff_method, grad_on_execution):
+    def test_classical_processing(
+        self,
+        interface,
+        dev,
+        diff_method,
+        grad_on_execution,
+        device_vjp,
+    ):
         """Test classical processing within the quantum tape"""
         a = torch.tensor(0.1, dtype=torch.float64, requires_grad=True)
         b = torch.tensor(0.2, dtype=torch.float64, requires_grad=False)
         c = torch.tensor(0.3, dtype=torch.float64, requires_grad=True)
 
         @qnode(
-            dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
+            dev,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            interface=interface,
+            device_vjp=device_vjp,
         )
         def circuit(a, b, c):
             qml.RY(a * c, wires=0)
@@ -336,11 +392,22 @@ class TestQNode:
         assert b.grad is None
         assert isinstance(c.grad, torch.Tensor)
 
-    def test_no_trainable_parameters(self, interface, dev, diff_method, grad_on_execution):
+    def test_no_trainable_parameters(
+        self,
+        interface,
+        dev,
+        diff_method,
+        grad_on_execution,
+        device_vjp,
+    ):
         """Test evaluation and Jacobian if there are no trainable parameters"""
 
         @qnode(
-            dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
+            dev,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            interface=interface,
+            device_vjp=device_vjp,
         )
         def circuit(a, b):
             qml.RY(a, wires=0)
@@ -384,14 +451,20 @@ class TestQNode:
             np.array([[0, 1], [1, 0]]),
         ],
     )
-    def test_matrix_parameter(self, interface, dev, diff_method, grad_on_execution, U, tol):
+    def test_matrix_parameter(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, U, tol
+    ):
         """Test that the Torch interface works correctly
         with a matrix parameter"""
         a_val = 0.1
         a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True)
 
         @qnode(
-            dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
+            dev,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            interface=interface,
+            device_vjp=device_vjp,
         )
         def circuit(U, a):
             qml.QubitUnitary(U, wires=0)
@@ -408,11 +481,16 @@ class TestQNode:
         res.backward()
         assert np.allclose(a.grad, np.sin(a_val), atol=tol, rtol=0)
 
-    def test_differentiable_expand(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_differentiable_expand(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test that operation and nested tapes expansion
         is differentiable"""
         kwargs = dict(
-            diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            interface=interface,
+            device_vjp=device_vjp,
         )
         if diff_method == "spsa":
             kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA)
@@ -561,12 +639,15 @@ class TestShotsIntegration:
 
 
 @pytest.mark.parametrize(
-    "interface,dev,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method
+    "interface,dev,diff_method,grad_on_execution, device_vjp",
+    interface_and_qubit_device_and_diff_method,
 )
 class TestQubitIntegration:
     """Tests that ensure various qubit circuits integrate correctly"""
 
-    def test_probability_differentiation(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_probability_differentiation(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
         kwargs = {}
@@ -586,6 +667,7 @@ class TestQubitIntegration:
             diff_method=diff_method,
             grad_on_execution=grad_on_execution,
             interface=interface,
+            device_vjp=device_vjp,
             **kwargs,
         )
         def circuit(x, y):
@@ -621,11 +703,16 @@ class TestQubitIntegration:
         assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0)
         assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0)
 
-    def test_ragged_differentiation(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_ragged_differentiation(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
         kwargs = dict(
-            diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            interface=interface,
+            device_vjp=device_vjp,
         )
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not currently support returning probabilities")
@@ -671,11 +758,22 @@ class TestQubitIntegration:
         assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0)
         assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0)
 
-    def test_chained_qnodes(self, interface, dev, diff_method, grad_on_execution):
+    def test_chained_qnodes(
+        self,
+        interface,
+        dev,
+        diff_method,
+        grad_on_execution,
+        device_vjp,
+    ):
         """Test that the gradient of chained QNodes works without error"""
 
         @qnode(
-            dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution
+            dev,
+            interface=interface,
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit1(weights):
             qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1])
@@ -707,7 +805,7 @@ class TestQubitIntegration:
         loss = cost(weights)
         loss.backward()
 
-    def test_hessian(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test hessian calculation of a scalar valued QNode"""
         if diff_method in {"adjoint", "spsa"}:
             pytest.skip("Adjoint and SPSA do not support second derivative.")
@@ -722,6 +820,7 @@ class TestQubitIntegration:
             grad_on_execution=grad_on_execution,
             max_diff=2,
             interface=interface,
+            device_vjp=device_vjp,
             **options,
         )
         def circuit(x):
@@ -757,7 +856,9 @@ class TestQubitIntegration:
         else:
             assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)
 
-    def test_hessian_vector_valued(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_hessian_vector_valued(
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test hessian calculation of a vector valued QNode"""
         if diff_method in {"adjoint", "spsa"}:
             pytest.skip("Adjoint and SPSA do not support second derivative.")
@@ -772,6 +873,7 @@ class TestQubitIntegration:
             grad_on_execution=grad_on_execution,
             max_diff=2,
             interface=interface,
+            device_vjp=device_vjp,
             **options,
         )
         def circuit(x):
@@ -818,7 +920,7 @@ class TestQubitIntegration:
         else:
             assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)
 
-    def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test hessian calculation of a ragged QNode"""
         if diff_method in {"adjoint", "spsa"}:
             pytest.skip("Adjoint and SPSA do not support second derivative.")
@@ -833,6 +935,7 @@ class TestQubitIntegration:
             grad_on_execution=grad_on_execution,
             max_diff=2,
             interface=interface,
+            device_vjp=device_vjp,
             **options,
         )
         def circuit(x):
@@ -891,7 +994,7 @@ class TestQubitIntegration:
             assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)
 
     def test_hessian_vector_valued_postprocessing(
-        self, interface, dev, diff_method, grad_on_execution, tol
+        self, interface, dev, diff_method, grad_on_execution, device_vjp, tol
     ):
         """Test hessian calculation of a vector valued QNode with post-processing"""
         if diff_method in {"adjoint", "spsa"}:
@@ -907,6 +1010,7 @@ class TestQubitIntegration:
             grad_on_execution=grad_on_execution,
             max_diff=2,
             interface=interface,
+            device_vjp=device_vjp,
             **options,
         )
         def circuit(x):
@@ -951,7 +1055,7 @@ class TestQubitIntegration:
         else:
             assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)
 
-    def test_state(self, interface, dev, diff_method, grad_on_execution, tol):
+    def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol):
         """Test that the state can be returned and differentiated"""
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support states")
@@ -960,7 +1064,11 @@ class TestQubitIntegration:
         y = torch.tensor(-0.654, requires_grad=True)
 
         @qnode(
-            dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            dev,
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -987,10 +1095,15 @@ class TestQubitIntegration:
         assert torch.allclose(res, expected, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("state", [[1], [0, 1]])  # Basis state and state vector
-    def test_projector(self, state, interface, dev, diff_method, grad_on_execution, tol):
+    def test_projector(
+        self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol
+    ):
         """Test that the variance of a projector is correctly returned"""
         kwargs = dict(
-            diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
+            diff_method=diff_method,
+            interface=interface,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         if diff_method == "adjoint":
             pytest.skip("Adjoint does not support projectors")
@@ -1029,12 +1142,14 @@ class TestQubitIntegration:
         assert np.allclose(weights.grad.detach(), expected, atol=tol, rtol=0)
 
 
-@pytest.mark.parametrize("dev,diff_method,grad_on_execution", qubit_device_and_diff_method)
+@pytest.mark.parametrize(
+    "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method
+)
 class TestTapeExpansion:
     """Test that tape expansion within the QNode integrates correctly
     with the Torch interface"""
 
-    def test_gradient_expansion(self, dev, diff_method, grad_on_execution):
+    def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp):
         """Test that a *supported* operation with no gradient recipe is
         expanded for both parameter-shift and finite-differences, but not for execution."""
         if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"):
@@ -1051,6 +1166,7 @@ class TestTapeExpansion:
             diff_method=diff_method,
             grad_on_execution=grad_on_execution,
             max_diff=2,
+            device_vjp=device_vjp,
             interface="torch",
         )
         def circuit(x):
@@ -1078,6 +1194,7 @@ class TestTapeExpansion:
         dev,
         diff_method,
         grad_on_execution,
+        device_vjp,
         max_diff,
     ):
         """Test that a *supported* operation with no gradient recipe is only
@@ -1097,6 +1214,7 @@ class TestTapeExpansion:
             grad_on_execution=grad_on_execution,
             max_diff=max_diff,
             interface="torch",
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.Hadamard(wires=0)
@@ -1112,7 +1230,7 @@ class TestTapeExpansion:
 
     @pytest.mark.parametrize("max_diff", [1, 2])
     def test_hamiltonian_expansion_analytic(
-        self, dev, diff_method, grad_on_execution, max_diff, tol
+        self, dev, diff_method, grad_on_execution, max_diff, device_vjp, tol
     ):
         """Test that if there
         are non-commuting groups and the number of shots is None
@@ -1122,6 +1240,7 @@ class TestTapeExpansion:
             grad_on_execution=grad_on_execution,
             max_diff=max_diff,
             interface="torch",
+            device_vjp=device_vjp,
         )
         if diff_method == "adjoint":
             pytest.skip("The adjoint method does not yet support Hamiltonians")
@@ -1191,7 +1310,9 @@ class TestTapeExpansion:
             assert torch.allclose(grad2_w_c, expected, atol=tol)
 
     @pytest.mark.parametrize("max_diff", [1, 2])
-    def test_hamiltonian_finite_shots(self, dev, diff_method, grad_on_execution, max_diff):
+    def test_hamiltonian_finite_shots(
+        self, dev, diff_method, device_vjp, grad_on_execution, max_diff
+    ):
         """Test that the Hamiltonian is correctly measured if there
         are non-commuting groups and the number of shots is finite
         and the first and second order gradients are correctly evaluated"""
@@ -1220,6 +1341,7 @@ class TestTapeExpansion:
             grad_on_execution=grad_on_execution,
             max_diff=max_diff,
             interface="torch",
+            device_vjp=device_vjp,
             **gradient_kwargs,
         )
         def circuit(data, weights, coeffs):
@@ -1400,17 +1522,20 @@ class TestSample:
 
 
 qubit_device_and_diff_method_and_grad_on_execution = [
-    [DefaultQubit(), "backprop", True],
-    [DefaultQubit(), "finite-diff", False],
-    [DefaultQubit(), "parameter-shift", False],
-    [DefaultQubit(), "adjoint", True],
-    [DefaultQubit(), "adjoint", False],
-    [DefaultQubit(), "hadamard", False],
+    [DefaultQubit(), "backprop", True, False],
+    [DefaultQubit(), "finite-diff", False, False],
+    [DefaultQubit(), "parameter-shift", False, False],
+    [DefaultQubit(), "adjoint", True, False],
+    [DefaultQubit(), "adjoint", False, False],
+    [DefaultQubit(), "adjoint", True, True],
+    [DefaultQubit(), "adjoint", False, True],
+    [DefaultQubit(), "hadamard", False, False],
 ]
 
 
 @pytest.mark.parametrize(
-    "dev,diff_method,grad_on_execution", qubit_device_and_diff_method_and_grad_on_execution
+    "dev,diff_method,grad_on_execution, device_vjp",
+    qubit_device_and_diff_method_and_grad_on_execution,
 )
 @pytest.mark.parametrize("shots", [None, 10000])
 class TestReturn:
@@ -1418,12 +1543,20 @@ class TestReturn:
 
     # pylint:disable=too-many-public-methods
 
-    def test_grad_single_measurement_param(self, dev, diff_method, grad_on_execution, shots):
+    def test_grad_single_measurement_param(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
+    ):
         """For one measurement and one param, the gradient is a float."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
+        )
         def circuit(a):
             qml.RY(a, wires=0)
             qml.RX(0.2, wires=0)
@@ -1443,13 +1576,19 @@ class TestReturn:
         assert grad.shape == ()
 
     def test_grad_single_measurement_multiple_param(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, shots, device_vjp
     ):
         """For one measurement and multiple param, the gradient is a tuple of arrays."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
+        )
         def circuit(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=0)
@@ -1469,13 +1608,19 @@ class TestReturn:
         assert grad_b.shape == ()
 
     def test_grad_single_measurement_multiple_param_array(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """For one measurement and multiple param as a single array params, the gradient is an array."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a):
             qml.RY(a[0], wires=0)
             qml.RX(a[1], wires=0)
@@ -1493,7 +1638,7 @@ class TestReturn:
         assert grad.shape == (2,)
 
     def test_jacobian_single_measurement_param_probs(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """For a multi dimensional measurement (probs), check that a single array is returned with the correct
         dimension"""
@@ -1503,7 +1648,13 @@ class TestReturn:
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of probabilities.")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a):
             qml.RY(a, wires=0)
             qml.RX(0.2, wires=0)
@@ -1517,7 +1668,7 @@ class TestReturn:
         assert jac.shape == (4,)
 
     def test_jacobian_single_measurement_probs_multiple_param(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
         the correct dimension"""
@@ -1526,7 +1677,13 @@ class TestReturn:
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=0)
@@ -1546,7 +1703,7 @@ class TestReturn:
         assert jac[1].shape == (4,)
 
     def test_jacobian_single_measurement_probs_multiple_param_single_array(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with
         the correct dimension"""
@@ -1555,7 +1712,13 @@ class TestReturn:
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a):
             qml.RY(a[0], wires=0)
             qml.RX(a[1], wires=0)
@@ -1568,7 +1731,7 @@ class TestReturn:
         assert jac.shape == (4, 2)
 
     def test_jacobian_expval_expval_multiple_params(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
@@ -1582,6 +1745,7 @@ class TestReturn:
             interface="torch",
             diff_method=diff_method,
             max_diff=1,
+            device_vjp=device_vjp,
             grad_on_execution=grad_on_execution,
         )
         def circuit(x, y):
@@ -1609,13 +1773,19 @@ class TestReturn:
         assert jac[1][1].shape == ()
 
     def test_jacobian_expval_expval_multiple_params_array(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a):
             qml.RY(a[0], wires=0)
             qml.RX(a[1], wires=0)
@@ -1634,7 +1804,9 @@ class TestReturn:
         assert isinstance(jac[1], torch.Tensor)
         assert jac[1].shape == (2,)
 
-    def test_jacobian_var_var_multiple_params(self, dev, diff_method, grad_on_execution, shots):
+    def test_jacobian_var_var_multiple_params(
+        self, dev, diff_method, device_vjp, grad_on_execution, shots
+    ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of var.")
@@ -1650,6 +1822,7 @@ class TestReturn:
             dev,
             interface="torch",
             diff_method=diff_method,
+            device_vjp=device_vjp,
             max_diff=1,
             grad_on_execution=grad_on_execution,
         )
@@ -1679,7 +1852,7 @@ class TestReturn:
         assert jac[1][1].shape == ()
 
     def test_jacobian_var_var_multiple_params_array(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, device_vjp, grad_on_execution, shots
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
         if diff_method == "adjoint":
@@ -1689,7 +1862,13 @@ class TestReturn:
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a):
             qml.RY(a[0], wires=0)
             qml.RX(a[1], wires=0)
@@ -1709,7 +1888,7 @@ class TestReturn:
         assert jac[1].shape == (2,)
 
     def test_jacobian_multiple_measurement_single_param(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """The jacobian of multiple measurements with a single params return an array."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
@@ -1718,7 +1897,13 @@ class TestReturn:
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because of probabilities.")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a):
             qml.RY(a, wires=0)
             qml.RX(0.2, wires=0)
@@ -1738,7 +1923,7 @@ class TestReturn:
         assert jac[1].shape == (4,)
 
     def test_jacobian_multiple_measurement_multiple_param(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """The jacobian of multiple measurements with a multiple params return a tuple of arrays."""
         if diff_method == "adjoint":
@@ -1746,7 +1931,13 @@ class TestReturn:
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            device_vjp=device_vjp,
+            grad_on_execution=grad_on_execution,
+        )
         def circuit(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=0)
@@ -1775,7 +1966,7 @@ class TestReturn:
         assert jac[1][1].shape == (4,)
 
     def test_jacobian_multiple_measurement_multiple_param_array(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """The jacobian of multiple measurements with a multiple params array return a single array."""
         if diff_method == "adjoint":
@@ -1783,7 +1974,13 @@ class TestReturn:
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
 
-        @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution)
+        @qnode(
+            dev,
+            interface="torch",
+            diff_method=diff_method,
+            grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
+        )
         def circuit(a):
             qml.RY(a[0], wires=0)
             qml.RX(a[1], wires=0)
@@ -1802,7 +1999,9 @@ class TestReturn:
         assert isinstance(jac[1], torch.Tensor)
         assert jac[1].shape == (4, 2)
 
-    def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_execution, shots):
+    def test_hessian_expval_multiple_params(
+        self, dev, diff_method, grad_on_execution, shots, device_vjp
+    ):
         """The hessian of single a measurement with multiple params return a tuple of arrays."""
         if shots is not None and diff_method in ("backprop", "adjoint"):
             pytest.skip("Test does not support finite shots and adjoint/backprop")
@@ -1819,6 +2018,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1845,7 +2045,9 @@ class TestReturn:
         assert hess[1][0].shape == ()
         assert hess[1][1].shape == ()
 
-    def test_hessian_expval_multiple_param_array(self, dev, diff_method, grad_on_execution, shots):
+    def test_hessian_expval_multiple_param_array(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
+    ):
         """The hessian of single measurement with a multiple params array return a single array."""
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because second order diff.")
@@ -1860,6 +2062,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -1872,7 +2075,9 @@ class TestReturn:
         assert isinstance(hess, torch.Tensor)
         assert hess.shape == (2, 2)
 
-    def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, shots):
+    def test_hessian_var_multiple_params(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
+    ):
         """The hessian of a single measurement with multiple params returns a tuple of arrays."""
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because second order diff.")
@@ -1890,6 +2095,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1916,7 +2122,9 @@ class TestReturn:
         assert isinstance(hess[1][1], torch.Tensor)
         assert hess[1][1].shape == ()
 
-    def test_hessian_var_multiple_param_array(self, dev, diff_method, grad_on_execution, shots):
+    def test_hessian_var_multiple_param_array(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
+    ):
         """The hessian of single measurement with a multiple params array return a single array."""
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because second order diff.")
@@ -1933,6 +2141,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -1945,7 +2154,9 @@ class TestReturn:
         assert isinstance(hess, torch.Tensor)
         assert hess.shape == (2, 2)
 
-    def test_hessian_probs_expval_multiple_params(self, dev, diff_method, grad_on_execution, shots):
+    def test_hessian_probs_expval_multiple_params(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
+    ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because second order diff.")
@@ -1964,6 +2175,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -1996,7 +2208,7 @@ class TestReturn:
         assert tuple(hess[1][1].shape) == (3,)
 
     def test_hessian_expval_probs_multiple_param_array(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """The hessian of multiple measurements with a multiple param array return a single array."""
         if diff_method == "adjoint":
@@ -2015,6 +2227,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])
@@ -2032,7 +2245,9 @@ class TestReturn:
         assert isinstance(hess, torch.Tensor)
         assert tuple(hess.shape) == (3, 2, 2)
 
-    def test_hessian_probs_var_multiple_params(self, dev, diff_method, grad_on_execution, shots):
+    def test_hessian_probs_var_multiple_params(
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
+    ):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
         if diff_method == "adjoint":
             pytest.skip("Test does not supports adjoint because second order diff.")
@@ -2051,6 +2266,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x, y):
             qml.RX(x, wires=[0])
@@ -2083,7 +2299,7 @@ class TestReturn:
         assert tuple(hess[1][1].shape) == (3,)
 
     def test_hessian_var_probs_multiple_param_array(
-        self, dev, diff_method, grad_on_execution, shots
+        self, dev, diff_method, grad_on_execution, device_vjp, shots
     ):
         """The hessian of multiple measurements with a multiple param array return a single array."""
         if diff_method == "adjoint":
@@ -2102,6 +2318,7 @@ class TestReturn:
             diff_method=diff_method,
             max_diff=2,
             grad_on_execution=grad_on_execution,
+            device_vjp=device_vjp,
         )
         def circuit(x):
             qml.RX(x[0], wires=[0])