diff --git a/pennylane/interfaces/autograd.py b/pennylane/interfaces/autograd.py index 7510b0036bb..c5465d2e283 100644 --- a/pennylane/interfaces/autograd.py +++ b/pennylane/interfaces/autograd.py @@ -118,7 +118,7 @@ def _execute( if isinstance(r[0][0], dict): # This happens when measurement type is Counts and shot vector is passed continue - except (IndexError, KeyError): + except (TypeError, IndexError, KeyError): pass r = np.hstack(r) if r.dtype == np.dtype("object") else r res[i] = np.tensor(r) diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py index 8d89861726a..a492d87dbba 100644 --- a/tests/devices/test_custom_return_obj.py +++ b/tests/devices/test_custom_return_obj.py @@ -62,6 +62,14 @@ class DeviceSupporingSpecialObservable(DefaultQubit): short_name = "default.qubit.specialobservable" observables = DefaultQubit.observables.union({"SpecialObservable"}) + @classmethod + def capabilities(cls): + capabilities = super().capabilities().copy() + capabilities.update( + provides_jacobian=True, + ) + return capabilities + def expval(self, observable, **kwargs): if self.analytic and isinstance(observable, SpecialObservable): val = super().expval(qml.PauliZ(wires=0), **kwargs) @@ -69,6 +77,13 @@ def expval(self, observable, **kwargs): return super().expval(observable, **kwargs) + def jacobian(self, tape): + # we actually let pennylane do the work of computing the + # jacobian for us but return it as a device jacobian + gradient_tapes, fn = qml.gradients.param_shift(tape) + tape_jacobian = fn(qml.execute(gradient_tapes, self, None)) + return tape_jacobian + dev = DeviceSupporingSpecialObservable(wires=1, shots=None) # force diff_method='parameter-shift' because otherwise @@ -87,7 +102,20 @@ def reference_qnode(x): assert isinstance(out, np.ndarray) assert isinstance(out.item(), SpecialObject) assert np.isclose(out.item().val, reference_qnode(0.2)) + reference_jac = qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)), + assert np.isclose( + reference_jac, qml.jacobian(qnode)(np.array(0.2, requires_grad=True)).item().val, - qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)), + ) + + # now check that also the device jacobian works with a custom return type + @qml.qnode(dev, diff_method="device") + def device_gradient_qnode(x): + qml.RY(x, wires=0) + return qml.expval(SpecialObservable(wires=0)) + + assert np.isclose( + reference_jac, + qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val, )