diff --git a/doc/releases/changelog-0.23.0.md b/doc/releases/changelog-0.23.0.md
index 0e82ad1a95b..268f34a96e2 100644
--- a/doc/releases/changelog-0.23.0.md
+++ b/doc/releases/changelog-0.23.0.md
@@ -614,12 +614,15 @@ the `decimals` and `show_matrices` keywords are added. `qml.drawer.tape_text(tap
the qubit operator that is prepared for tapering the HF state.
[(#2441)](https://github.com/PennyLaneAI/pennylane/pull/2441)
+* Fixed a bug with custom device defined jacobians not being returned properly.
+ [(#2485)](https://github.com/PennyLaneAI/pennylane-sf/pull/2485)
+
Documentation
Contributors
This release contains contributions from (in alphabetical order):
-Karim Alaa El-Din, Guillermo Alonso-Linaje, Juan Miguel Arrazola, Ali Asadi, Utkarsh Azad, Thomas Bromley, Alain Delgado,
-Olivia Di Matteo, Anthony Hayes, David Ittah, Josh Izaac, Soran Jahangiri, Christina Lee, Romain Moyard, Zeyue Niu,
-Matthew Silverman, Lee James O'Riordan, Jay Soni, Antal Száva, Maurice Weber, David Wierichs.
+Karim Alaa El-Din, Guillermo Alonso-Linaje, Juan Miguel Arrazola, Ali Asadi, Utkarsh Azad, Samuel Banning,
+Thomas Bromley, Alain Delgado, Olivia Di Matteo, Anthony Hayes, David Ittah, Josh Izaac, Soran Jahangiri, Christina Lee,
+Romain Moyard, Zeyue Niu, Matthew Silverman, Lee James O'Riordan, Jay Soni, Antal Száva, Maurice Weber, David Wierichs.
diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py
index 9053492a4c9..d09aebb01cc 100644
--- a/pennylane/interfaces/autograd.py
+++ b/pennylane/interfaces/autograd.py
@@ -229,7 +229,15 @@ def grad_fn(dy):
vjps = [qml.gradients.compute_vjp(d, jac) for d, jac in zip(dy, jacs)]
- return [qml.math.to_numpy(v, max_depth=_n) if isinstance(v, ArrayBox) else v for v in vjps]
+ return_vjps = [
+ qml.math.to_numpy(v, max_depth=_n) if isinstance(v, ArrayBox) else v for v in vjps
+ ]
+ if device.capabilities().get("provides_jacobian", False):
+ # in the case where the device provides the jacobian,
+ # the output of grad_fn must be wrapped in a tuple in
+ # order to match the input parameters to _execute.
+ return (return_vjps,)
+ return return_vjps
return grad_fn
diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index c20ef93c8b9..994d46f550c 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -21,6 +21,7 @@
from pennylane import numpy as np
import pennylane as qml
+from pennylane.devices import DefaultQubit
from pennylane.gradients import finite_diff, param_shift
from pennylane.interfaces import execute
@@ -1111,3 +1112,30 @@ def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol):
res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev))
expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)
assert np.allclose(res, expected, atol=tol, rtol=0)
+
+
+class TestCustomJacobian:
+ def test_custom_jacobians(self):
+ class CustomJacobianDevice(DefaultQubit):
+ @classmethod
+ def capabilities(cls):
+ capabilities = super().capabilities()
+ capabilities["provides_jacobian"] = True
+ return capabilities
+
+ def jacobian(self, tape):
+ return np.array([1.0, 2.0, 3.0, 4.0])
+
+ dev = CustomJacobianDevice(wires=2)
+
+ @qml.qnode(dev, diff_method="device")
+ def circuit(v):
+ qml.RX(v, wires=0)
+ return qml.probs(wires=[0, 1])
+
+ d_circuit = qml.jacobian(circuit, argnum=0)
+
+ params = np.array(1.0, requires_grad=True)
+
+ d_out = d_circuit(params)
+ assert np.allclose(d_out, np.array([1.0, 2.0, 3.0, 4.0]))