Compatibility with the new return system

pennylane_qiskit/qiskit_device.py

@@ -22,6 +22,7 @@
 import warnings
 
 import numpy as np
+import pennylane
 from qiskit import ClassicalRegister, QuantumCircuit, QuantumRegister
 from qiskit import extensions as ex
 from qiskit.compiler import transpile
@@ -479,9 +480,15 @@ def batch_execute(self, circuits):
             if self.shots is not None or circuit.is_sampled:
                 self._samples = self.generate_samples(circuit_obj)
 
-            res = self.statistics(circuit)
-            res = np.asarray(res)
-            results.append(res)
+            if not pennylane.active_return():
+                res = self._statistics_legacy(circuit)
+                res = np.asarray(res)
+                results.append(res)
+            else:
+                res = self.statistics(circuit)
+                single_measurement = len(circuit.measurements) == 1
+                res = res[0] if single_measurement else tuple(res)
+                results.append(res)
 
         if self.tracker.active:
             self.tracker.update(batches=1, batch_len=len(circuits))

tests/test_converter.py

@@ -1115,21 +1115,14 @@ def circuit(a_val, b_val):
 
         assert np.allclose(res, res_expected)
 
-        jac = qml.jacobian(circuit)(x, y)
-
-        # Note: with v0.21.0, the behaviour of qml.jacobian changed. Branch
-        # based on the version number to ensure that the test case passes even
-        # with <v0.21.
-        if qml.__version__ < "0.21.0":
-            jac_expected = [
-                [-np.sin(x + np.cos(y)), np.sin(x + np.cos(y)) * np.sin(y)],
-                [np.cos(x * y) * y, np.cos(x * y) * x],
-            ]
-
-        else:
-            jac_expected = [
-                [-np.sin(x + np.cos(y)), np.cos(x * y) * y],
-                [np.sin(x + np.cos(y)) * np.sin(y), np.cos(x * y) * x],
-            ]
+        def cost(x, y):
+            return qml.math.stack(circuit(x, y))
+
+        jac = qml.jacobian(cost)(x, y)
+
+        jac_expected = [
+            [-np.sin(x + np.cos(y)), np.cos(x * y) * y],
+            [np.sin(x + np.cos(y)) * np.sin(y), np.cos(x * y) * x],
+        ]
 
         assert np.allclose(jac, jac_expected)

tests/test_qiskit_device.py

@@ -1,4 +1,5 @@
 import numpy as np
+import pennylane
 import pytest
 
 import pennylane as qml
@@ -62,14 +63,14 @@ def test_warning_raised_for_hardware_backend_analytic_expval(self, hardware_back
         assert len(record) == 1
         # check that the message matches
         assert (
-            record[0].message.args[0] == "The analytic calculation of "
-            "expectations, variances and probabilities is only supported on "
-            "statevector backends, not on the {}. Such statistics obtained from this "
-            "device are estimates based on samples.".format(dev.backend)
+                record[0].message.args[0] == "The analytic calculation of "
+                                             "expectations, variances and probabilities is only supported on "
+                                             "statevector backends, not on the {}. Such statistics obtained from this "
+                                             "device are estimates based on samples.".format(dev.backend)
         )
 
     def test_no_warning_raised_for_software_backend_analytic_expval(
-        self, statevector_backend, recorder, recwarn
+            self, statevector_backend, recorder, recwarn
     ):
         """Tests that no warning is raised if the analytic attribute is true on
         statevector simulators when calculating the expectation"""
@@ -140,8 +141,10 @@ def test_calls_to_reset(self, n_tapes, mocker, device):
 
         assert spy.call_count == n_tapes
 
-    def test_result(self, device, tol):
+    def test_result_legacy(self, device, tol):
         """Tests that the result has the correct shape and entry types."""
+        # TODO: remove once the legacy return system is activated.
+        pennylane.disable_return()
         dev = device(2)
         tapes = [self.tape1, self.tape2]
         res = dev.batch_execute(tapes)
@@ -161,6 +164,30 @@ def test_result(self, device, tol):
         assert isinstance(res[1], np.ndarray)
         assert np.allclose(res[1], tape2_expected, atol=0)
 
+        pennylane.enable_return()
+
+    def test_result(self, device, tol):
+        """Tests that the result has the correct shape and entry types."""
+        dev = device(2)
+        tapes = [self.tape1, self.tape2]
+        res = dev.batch_execute(tapes)
+
+        # We're calling device methods directly, need to reset before the next
+        # execution
+        dev.reset()
+        tape1_expected = dev.execute(self.tape1)
+
+        dev.reset()
+        tape2_expected = dev.execute(self.tape2)
+
+        assert len(res) == 2
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == 2
+        assert np.allclose(qml.math.stack(res[0]), tape1_expected, atol=0)
+
+        assert isinstance(res[1], np.ndarray)
+        assert np.allclose(res[1], tape2_expected, atol=0)
+
     def test_result_empty_tape(self, device, tol):
         """Tests that the result has the correct shape and entry types for
         empty tapes."""