added custom return object (#2808)

* added custom return object

* added comparison

* added more asserts

* progress

* changed description

* Update pennylane/_qubit_device.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update tests/devices/test_custom_return_obj.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Update tests/devices/test_custom_return_obj.py

Co-authored-by: antalszava <antalszava@gmail.com>

* Custom return type (#2813)

* fix problem from #2808 (comment)

* also test device gradient with custom return type

* Empty-Commit

* fixed formatting for test

* moved test to existing autograd test file

* Update doc/releases/changelog-dev.md

Co-authored-by: antalszava <antalszava@gmail.com>

* reverted

* fixed

* fixed change

* seperated into 2 testcases

* fixed black

Co-authored-by: antalszava <antalszava@gmail.com>
Co-authored-by: cvjjm <60615188+cvjjm@users.noreply.github.com>

doc/releases/changelog-dev.md

@@ -403,17 +403,20 @@ of operators. [(#2622)](https://github.com/PennyLaneAI/pennylane/pull/2622)
 * The adjoint of an adjoint has a correct `expand` result.
   [(#2766)](https://github.com/PennyLaneAI/pennylane/pull/2766)
 
+* Fix the ability to return custom objects as the expectation value of a QNode with the Autograd interface.
+  [(#2808)](https://github.com/PennyLaneAI/pennylane/pull/2808)
+
 * The WireCut operator now raises an error when instantiating it with an empty list.
   [(#2826)](https://github.com/PennyLaneAI/pennylane/pull/2826)
 
 * Allow hamiltonians with grouped observables to be measured on devices
   which were transformed using `qml.transform.insert()`.
-  [(#2857)](https://github.com/PennyLaneAI/pennylane/pull/2857) 
+  [(#2857)](https://github.com/PennyLaneAI/pennylane/pull/2857)
 
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
 
-Juan Miguel Arrazola, David Ittah, Soran Jahangiri, Edward Jiang, Ankit Khandelwal, Christina Lee,
+Samuel Banning, Juan Miguel Arrazola, David Ittah, Soran Jahangiri, Edward Jiang, Ankit Khandelwal, Christina Lee,
 Sergio Martínez-Losa, Albert Mitjans Coma, Ixchel Meza Chavez, Romain Moyard, Lee James O'Riordan,
 Mudit Pandey, Bogdan Reznychenko, Jay Soni, Antal Száva, David Wierichs, Moritz Willmann

pennylane/_qubit_device.py

@@ -332,7 +332,11 @@ def execute(self, circuit, **kwargs):
                     results = self._asarray(results, dtype=self.C_DTYPE)
                 elif circuit.measurements[0].return_type is not qml.measurements.Counts:
                     # Measurements with expval, var or probs
-                    results = self._asarray(results, dtype=self.R_DTYPE)
+                    try:
+                        # Feature for returning custom objects: if the type cannot be cast to float then we can still allow it as an output
+                        results = self._asarray(results, dtype=self.R_DTYPE)
+                    except TypeError:
+                        pass
 
             elif all(
                 ret in (qml.measurements.Expectation, qml.measurements.Variance)

tests/interfaces/test_autograd.py

@@ -19,6 +19,7 @@
 import autograd
 import pytest
 from pennylane import numpy as np
+from pennylane.operation import Observable, AnyWires
 
 import pennylane as qml
 from pennylane.devices import DefaultQubit
@@ -1179,3 +1180,141 @@ def qnode(a, b):
 
         assert np.allclose(res[0], expected[0])
         assert np.allclose(res[1], expected[1])
+
+
+class SpecialObject:
+    """SpecialObject
+    A special object that conveniently encapsulates the return value of
+    a special observable supported by a special device and which supports
+    multiplication with scalars and addition.
+    """
+
+    def __init__(self, val):
+        self.val = val
+
+    def __mul__(self, other):
+        new = SpecialObject(self.val)
+        new *= other
+        return new
+
+    def __imul__(self, other):
+        self.val *= other
+        return self
+
+    def __rmul__(self, other):
+        return self * other
+
+    def __iadd__(self, other):
+        self.val += other.val if isinstance(other, self.__class__) else other
+        return self
+
+    def __add__(self, other):
+        new = SpecialObject(self.val)
+        new += other.val if isinstance(other, self.__class__) else other
+        return new
+
+    def __radd__(self, other):
+        return self + other
+
+
+class SpecialObservable(Observable):
+    """SpecialObservable"""
+
+    num_wires = AnyWires
+    num_params = 0
+    par_domain = None
+
+    def diagonalizing_gates(self):
+        """Diagonalizing gates"""
+        return []
+
+
+class DeviceSupportingSpecialObservable(DefaultQubit):
+    name = "Device supporting SpecialObservable"
+    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)
+            return SpecialObject(val)
+
+        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
+
+
+@pytest.mark.autograd
+class TestObservableWithObjectReturnType:
+    """Unit tests for qnode returning a custom object"""
+
+    def test_custom_return_type(self):
+        """Test custom return values for a qnode"""
+
+        dev = DeviceSupportingSpecialObservable(wires=1, shots=None)
+
+        # force diff_method='parameter-shift' because otherwise
+        # PennyLane swaps out dev for default.qubit.autograd
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def qnode(x):
+            qml.RY(x, wires=0)
+            return qml.expval(SpecialObservable(wires=0))
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def reference_qnode(x):
+            qml.RY(x, wires=0)
+            return qml.expval(qml.PauliZ(wires=0))
+
+        out = qnode(0.2)
+        assert isinstance(out, np.ndarray)
+        assert isinstance(out.item(), SpecialObject)
+        assert np.isclose(out.item().val, reference_qnode(0.2))
+
+    def test_jacobian_with_custom_return_type(self):
+        """Test differentiation of a QNode on a device supporting a
+        special observable that returns an object rather than a number."""
+
+        dev = DeviceSupportingSpecialObservable(wires=1, shots=None)
+
+        # force diff_method='parameter-shift' because otherwise
+        # PennyLane swaps out dev for default.qubit.autograd
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def qnode(x):
+            qml.RY(x, wires=0)
+            return qml.expval(SpecialObservable(wires=0))
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def reference_qnode(x):
+            qml.RY(x, wires=0)
+            return qml.expval(qml.PauliZ(wires=0))
+
+        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,
+        )
+
+        # 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,
+        )