Jacobian Slice Bug StrawberryFields

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)
+
 <h3>Documentation</h3>
 
 <h3>Contributors</h3>
 
 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.

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
 

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