From a998cc219daed9e0a38917a28c746e55cd68de66 Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 11:25:52 -0400
Subject: [PATCH 01/18] added custom return object

---
 pennylane/_qubit_device.py              |  5 +-
 tests/devices/test_custom_return_obj.py | 84 +++++++++++++++++++++++++
 2 files changed, 88 insertions(+), 1 deletion(-)
 create mode 100644 tests/devices/test_custom_return_obj.py

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 6dc8a01bc06..0da3c476f4a 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -309,7 +309,10 @@ 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:
+                        results = self._asarray(results, dtype=self.R_DTYPE)
+                    except TypeError:
+                        pass
 
             elif all(
                 ret in (qml.measurements.Expectation, qml.measurements.Variance)
diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
new file mode 100644
index 00000000000..77c9a362d35
--- /dev/null
+++ b/tests/devices/test_custom_return_obj.py
@@ -0,0 +1,84 @@
+import pennylane as qml
+from pennylane.devices import DefaultQubit
+from pennylane.operation import Observable, AnyWires
+
+
+class TestObservableWithObjectReturnType:
+    """Unit tests for differentiation of observables returning an object"""
+
+    def test_custom_return_type(self):
+        """Test differentiation of a QNode on a device supporting a
+        special observable that returns an object rathern than a nummber."""
+
+        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 DeviceSupporingSpecialObservable(DefaultQubit):
+            name = "Device supporing SpecialObservable"
+            short_name = "default.qibit.specialobservable"
+            observables = DefaultQubit.observables.union({"SpecialObservable"})
+
+            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)
+
+        dev = DeviceSupporingSpecialObservable(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))
+
+        assert isinstance(qnode(0.2), SpecialObject)

From d33310015fdab376a66c3a27d8cc7d198aa5ce32 Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 11:33:08 -0400
Subject: [PATCH 02/18] added comparison

---
 tests/devices/test_custom_return_obj.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
index 77c9a362d35..158a0287bd7 100644
--- a/tests/devices/test_custom_return_obj.py
+++ b/tests/devices/test_custom_return_obj.py
@@ -2,6 +2,8 @@
 from pennylane.devices import DefaultQubit
 from pennylane.operation import Observable, AnyWires
 
+import pennylane.numpy as np
+
 
 class TestObservableWithObjectReturnType:
     """Unit tests for differentiation of observables returning an object"""
@@ -82,3 +84,4 @@ def reference_qnode(x):
             return qml.expval(qml.PauliZ(wires=0))
 
         assert isinstance(qnode(0.2), SpecialObject)
+        assert np.isclose(qnode(0.2).item().val, reference_qnode(0.2))

From 03ddd79118be19febb2acbcd9946b3676693eb04 Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 11:39:50 -0400
Subject: [PATCH 03/18] added more asserts

---
 tests/devices/test_custom_return_obj.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
index 158a0287bd7..01c26e5aab5 100644
--- a/tests/devices/test_custom_return_obj.py
+++ b/tests/devices/test_custom_return_obj.py
@@ -83,5 +83,11 @@ def reference_qnode(x):
             qml.RY(x, wires=0)
             return qml.expval(qml.PauliZ(wires=0))
 
-        assert isinstance(qnode(0.2), SpecialObject)
-        assert np.isclose(qnode(0.2).item().val, reference_qnode(0.2))
+        out = qnode(0.2)
+        assert isinstance(out, np.ndarray)
+        assert isinstance(out.item(), SpecialObject)
+        assert np.isclose(out.item().val, reference_qnode(0.2))
+        assert np.isclose(
+            qml.jacobian(qnode)(np.array(0.2, requires_grad=True)).item().val,
+            qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)),
+        )

From 84559eb474b5972c334c7ef6a419fa2719cbaa4d Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 11:44:57 -0400
Subject: [PATCH 04/18] progress

---
 pennylane/_qubit_device.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 0da3c476f4a..79162d1c507 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -310,6 +310,7 @@ def execute(self, circuit, **kwargs):
                 elif circuit.measurements[0].return_type is not qml.measurements.Counts:
                     # Measurements with expval, var or probs
                     try:
+                        # experimental feature: if 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

From 7316d0872be2f0447b8b79c96d99510afd6cf5a5 Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 14:08:06 -0400
Subject: [PATCH 05/18] changed description

---
 tests/devices/test_custom_return_obj.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
index 01c26e5aab5..dfc49cd55e9 100644
--- a/tests/devices/test_custom_return_obj.py
+++ b/tests/devices/test_custom_return_obj.py
@@ -6,7 +6,7 @@
 
 
 class TestObservableWithObjectReturnType:
-    """Unit tests for differentiation of observables returning an object"""
+    """Unit tests for qnode returning a custom object"""
 
     def test_custom_return_type(self):
         """Test differentiation of a QNode on a device supporting a

From 5fb708f5faa74f4f38d4a8bd153e8b38e347503f Mon Sep 17 00:00:00 2001
From: Sam Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 16:24:00 -0400
Subject: [PATCH 06/18] Update pennylane/_qubit_device.py

Co-authored-by: antalszava <antalszava@gmail.com>
---
 pennylane/_qubit_device.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 79162d1c507..dcf7ed58474 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -310,7 +310,7 @@ def execute(self, circuit, **kwargs):
                 elif circuit.measurements[0].return_type is not qml.measurements.Counts:
                     # Measurements with expval, var or probs
                     try:
-                        # experimental feature: if type cannot be cast to float then we can still allow it as an output
+                        # 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

From 3bda482036fe76a8bada05f77f7ab0c2cb8d1356 Mon Sep 17 00:00:00 2001
From: Sam Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 16:24:20 -0400
Subject: [PATCH 07/18] Update tests/devices/test_custom_return_obj.py

Co-authored-by: antalszava <antalszava@gmail.com>
---
 tests/devices/test_custom_return_obj.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
index dfc49cd55e9..95bbdfdec4a 100644
--- a/tests/devices/test_custom_return_obj.py
+++ b/tests/devices/test_custom_return_obj.py
@@ -58,7 +58,7 @@ def diagonalizing_gates(self):
                 return []
 
         class DeviceSupporingSpecialObservable(DefaultQubit):
-            name = "Device supporing SpecialObservable"
+            name = "Device supporting SpecialObservable"
             short_name = "default.qibit.specialobservable"
             observables = DefaultQubit.observables.union({"SpecialObservable"})
 

From cca078514c3721bc4bc0be9adb712505235e2e85 Mon Sep 17 00:00:00 2001
From: Sam Banning <samcbanning@gmail.com>
Date: Wed, 13 Jul 2022 16:26:43 -0400
Subject: [PATCH 08/18] Update tests/devices/test_custom_return_obj.py

Co-authored-by: antalszava <antalszava@gmail.com>
---
 tests/devices/test_custom_return_obj.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
index 95bbdfdec4a..8d89861726a 100644
--- a/tests/devices/test_custom_return_obj.py
+++ b/tests/devices/test_custom_return_obj.py
@@ -59,7 +59,7 @@ def diagonalizing_gates(self):
 
         class DeviceSupporingSpecialObservable(DefaultQubit):
             name = "Device supporting SpecialObservable"
-            short_name = "default.qibit.specialobservable"
+            short_name = "default.qubit.specialobservable"
             observables = DefaultQubit.observables.union({"SpecialObservable"})
 
             def expval(self, observable, **kwargs):

From 6e9a17b2d95f474a22bb91e1e472bd2a01508eac Mon Sep 17 00:00:00 2001
From: cvjjm <60615188+cvjjm@users.noreply.github.com>
Date: Mon, 18 Jul 2022 15:33:42 +0200
Subject: [PATCH 09/18] Custom return type (#2813)

* fix problem from https://github.com/PennyLaneAI/pennylane/pull/2808#issuecomment-1184082711

* also test device gradient with custom return type
---
 pennylane/interfaces/autograd.py        |  2 +-
 tests/devices/test_custom_return_obj.py | 30 ++++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

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

From 68c656bcbbc5692b995f400b4ec53c82b977e93d Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Fri, 22 Jul 2022 10:27:38 -0400
Subject: [PATCH 10/18] Empty-Commit


From 18e92a909371c1fe991ba64c51d1b59a0cc16ddb Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Fri, 22 Jul 2022 14:16:27 -0400
Subject: [PATCH 11/18] fixed formatting for test

---
 tests/devices/test_custom_return_obj.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
index a492d87dbba..50806831cec 100644
--- a/tests/devices/test_custom_return_obj.py
+++ b/tests/devices/test_custom_return_obj.py
@@ -102,7 +102,7 @@ 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)),
+        reference_jac = (qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)),)
 
         assert np.isclose(
             reference_jac,

From 3d2dd35a2e28a9aee276030b964942a5e7ff59de Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Mon, 25 Jul 2022 11:06:37 -0400
Subject: [PATCH 12/18] moved test to existing autograd test file

---
 doc/releases/changelog-dev.md                |   5 +-
 tests/devices/test_custom_return_obj.py      | 121 -------------------
 tests/devices/test_default_qubit_autograd.py | 119 ++++++++++++++++++
 3 files changed, 123 insertions(+), 122 deletions(-)
 delete mode 100644 tests/devices/test_custom_return_obj.py

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 679dd4e7761..fb1e64602bd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -259,9 +259,12 @@
 * The adjoint of an adjoint has a correct `expand` result.
   [(#2766)](https://github.com/PennyLaneAI/pennylane/pull/2766)
 
+* Allow allow custom return types from autograd qnode.
+  [(#2808)](https://github.com/PennyLaneAI/pennylane/pull/2808)
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
 
-David Ittah, Edward Jiang, Ankit Khandelwal, Christina Lee, Sergio Martínez-Losa, Ixchel Meza Chavez, 
+Samuel Banning, David Ittah, Edward Jiang, Ankit Khandelwal, Christina Lee, Sergio Martínez-Losa, Ixchel Meza Chavez, 
 Lee James O'Riordan, Mudit Pandey, Bogdan Reznychenko, Jay Soni, Antal Száva, Moritz Willmann
\ No newline at end of file
diff --git a/tests/devices/test_custom_return_obj.py b/tests/devices/test_custom_return_obj.py
deleted file mode 100644
index 50806831cec..00000000000
--- a/tests/devices/test_custom_return_obj.py
+++ /dev/null
@@ -1,121 +0,0 @@
-import pennylane as qml
-from pennylane.devices import DefaultQubit
-from pennylane.operation import Observable, AnyWires
-
-import pennylane.numpy as np
-
-
-class TestObservableWithObjectReturnType:
-    """Unit tests for qnode returning a custom object"""
-
-    def test_custom_return_type(self):
-        """Test differentiation of a QNode on a device supporting a
-        special observable that returns an object rathern than a nummber."""
-
-        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 DeviceSupporingSpecialObservable(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
-
-        dev = DeviceSupporingSpecialObservable(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))
-        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,
-        )
diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py
index 79387676529..9d434b17307 100644
--- a/tests/devices/test_default_qubit_autograd.py
+++ b/tests/devices/test_default_qubit_autograd.py
@@ -19,6 +19,8 @@
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.devices.default_qubit_autograd import DefaultQubitAutograd
+from pennylane.devices import DefaultQubit
+from pennylane.operation import Observable, AnyWires
 from pennylane import DeviceError
 
 
@@ -541,3 +543,120 @@ def test_partial_subsystem(self, mocker):
 
         assert np.all(res == state)
         spy.assert_called()
+
+
+@pytest.mark.autograd
+class TestObservableWithObjectReturnType:
+    """Unit tests for qnode returning a custom object"""
+
+    def test_custom_return_type(self):
+        """Test differentiation of a QNode on a device supporting a
+        special observable that returns an object rathern than a nummber."""
+
+        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
+
+        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))
+        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,
+        )

From 166401e09b9c3237df178ba3c129a2553ef23faf Mon Sep 17 00:00:00 2001
From: Sam Banning <samcbanning@gmail.com>
Date: Mon, 25 Jul 2022 12:27:51 -0400
Subject: [PATCH 13/18] Update doc/releases/changelog-dev.md

Co-authored-by: antalszava <antalszava@gmail.com>
---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fb1e64602bd..3f7ea3f0c97 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -259,7 +259,7 @@
 * The adjoint of an adjoint has a correct `expand` result.
   [(#2766)](https://github.com/PennyLaneAI/pennylane/pull/2766)
 
-* Allow allow custom return types from autograd qnode.
+* 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)
 
 <h3>Contributors</h3>

From 430f0a26aae7b6a34885ac4a0e7385d074522705 Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Tue, 26 Jul 2022 10:23:40 -0400
Subject: [PATCH 14/18] reverted

---
 tests/devices/test_default_qubit_autograd.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py
index f35a805b529..79387676529 100644
--- a/tests/devices/test_default_qubit_autograd.py
+++ b/tests/devices/test_default_qubit_autograd.py
@@ -19,8 +19,6 @@
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.devices.default_qubit_autograd import DefaultQubitAutograd
-from pennylane.devices import DefaultQubit
-from pennylane.operation import Observable, AnyWires
 from pennylane import DeviceError
 
 
@@ -543,4 +541,3 @@ def test_partial_subsystem(self, mocker):
 
         assert np.all(res == state)
         spy.assert_called()
-

From 45c9da94e0dbc37bfc4e3a8a319c43bb7a71aaad Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Tue, 26 Jul 2022 10:24:41 -0400
Subject: [PATCH 15/18] fixed

---
 tests/interfaces/test_autograd.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index 5cf3ee42bee..639ca9c908f 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -18,7 +18,7 @@
 
 import autograd
 import pytest
-from pennylane import numpy as np
+from pennylane import numpy as np, Observable, AnyWires
 
 import pennylane as qml
 from pennylane.devices import DefaultQubit

From f9d7bf49c856a5421e6aca9753a6b4b2fc767f94 Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Tue, 26 Jul 2022 10:28:34 -0400
Subject: [PATCH 16/18] fixed change

---
 tests/interfaces/test_autograd.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index 639ca9c908f..7a044bfa37c 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -18,7 +18,8 @@
 
 import autograd
 import pytest
-from pennylane import numpy as np, Observable, AnyWires
+from pennylane import numpy as np
+from pennylane.operation import Observable, AnyWires
 
 import pennylane as qml
 from pennylane.devices import DefaultQubit

From 22b6dca1ecd38774e277a29a83b579bbc6ddd159 Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Tue, 26 Jul 2022 10:42:24 -0400
Subject: [PATCH 17/18] seperated into 2 testcases

---
 tests/interfaces/test_autograd.py | 165 +++++++++++++++++-------------
 1 file changed, 93 insertions(+), 72 deletions(-)

diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index 7a044bfa37c..c7e0f9cb33a 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -1182,85 +1182,110 @@ def qnode(a, b):
         assert np.allclose(res[1], expected[1])
 
 
-@pytest.mark.autograd
-class TestObservableWithObjectReturnType:
-    """Unit tests for qnode returning a custom object"""
-
-    def test_custom_return_type(self):
-        """Test differentiation of a QNode on a device supporting a
-        special observable that returns an object rathern than a nummber."""
-
-        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
+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 __imul__(self, other):
-                self.val *= other
-                return self
+    def expval(self, observable, **kwargs):
+        if self.analytic and isinstance(observable, SpecialObservable):
+            val = super().expval(qml.PauliZ(wires=0), **kwargs)
+            return SpecialObject(val)
 
-            def __rmul__(self, other):
-                return self * other
+        return super().expval(observable, **kwargs)
 
-            def __iadd__(self, other):
-                self.val += other.val if isinstance(other, self.__class__) else other
-                return self
+    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
 
-            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
+@pytest.mark.autograd
+class TestObservableWithObjectReturnType:
+    """Unit tests for qnode returning a custom object"""
 
-        class SpecialObservable(Observable):
-            """SpecialObservable"""
+    def test_custom_return_type(self):
+        """Test custom return values for a qnode"""
 
-            num_wires = AnyWires
-            num_params = 0
-            par_domain = None
+        dev = DeviceSupportingSpecialObservable(wires=1, shots=None)
 
-            def diagonalizing_gates(self):
-                """Diagonalizing gates"""
-                return []
+        # 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))
 
-        class DeviceSupportingSpecialObservable(DefaultQubit):
-            name = "Device supporting SpecialObservable"
-            short_name = "default.qubit.specialobservable"
-            observables = DefaultQubit.observables.union({"SpecialObservable"})
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def reference_qnode(x):
+            qml.RY(x, wires=0)
+            return qml.expval(qml.PauliZ(wires=0))
 
-            @classmethod
-            def capabilities(cls):
-                capabilities = super().capabilities().copy()
-                capabilities.update(
-                    provides_jacobian=True,
-                )
-                return capabilities
+        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 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
+    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)
 
@@ -1276,10 +1301,6 @@ 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))
         reference_jac = (qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)),)
 
         assert np.isclose(
@@ -1295,5 +1316,5 @@ def device_gradient_qnode(x):
 
         assert np.isclose(
             reference_jac,
-            qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val,
-        )
+            qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val
+        )
\ No newline at end of file

From ce921d8b20698e18bc025e686534d2fcd69d5dfa Mon Sep 17 00:00:00 2001
From: Samuel Banning <samcbanning@gmail.com>
Date: Tue, 26 Jul 2022 10:48:18 -0400
Subject: [PATCH 18/18] fixed black

---
 tests/interfaces/test_autograd.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index c7e0f9cb33a..396d0ba8646 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -1216,6 +1216,7 @@ def __add__(self, other):
     def __radd__(self, other):
         return self + other
 
+
 class SpecialObservable(Observable):
     """SpecialObservable"""
 
@@ -1227,6 +1228,7 @@ def diagonalizing_gates(self):
         """Diagonalizing gates"""
         return []
 
+
 class DeviceSupportingSpecialObservable(DefaultQubit):
     name = "Device supporting SpecialObservable"
     short_name = "default.qubit.specialobservable"
@@ -1281,8 +1283,6 @@ def reference_qnode(x):
         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."""
@@ -1316,5 +1316,5 @@ def device_gradient_qnode(x):
 
         assert np.isclose(
             reference_jac,
-            qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val
-        )
\ No newline at end of file
+            qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val,
+        )