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]))