Repository: PennyLaneAI/pennylane
Pull request #3315 (2022-11-10 20:06:57)

[Fix return-types] Remove backprop stack for Autograd and Tf 
**Changes:**

Remove the internal stacking when using backprop with autograd and tf.


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.

Patch:
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 3efa2c735..590f7b017 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -638,18 +638,6 @@ class QNode:
 
             res = res[0]
 
-            # Autograd or tensorflow: they do not support tuple return with backpropagation
-            backprop = False
-            if not isinstance(
-                self._qfunc_output, qml.measurements.MeasurementProcess
-            ) and self.interface in ("tf", "autograd"):
-                if isinstance(res, Sequence):
-                    backprop = any(qml.math.in_backprop(x) for x in res)
-                else:
-                    # res might not be a sequence even if the qfunc output is a sequence
-                    # of length 1
-                    backprop = qml.math.in_backprop(res)
-
             if old_interface == "auto":
                 self.interface = "auto"
 
@@ -657,14 +645,8 @@ class QNode:
             if isinstance(self._qfunc_output, list) and len(self._qfunc_output) == 1:
                 return [res]
 
-            if self.gradient_fn == "backprop" and backprop:
-                res = self.device._asarray(res)
-
             # If the return type is not tuple (list or ndarray) (Autograd and TF backprop removed)
-            if (
-                not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess))
-                and not backprop
-            ):
+            if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
                 if self.device._shot_vector:
                     res = [type(self.tape._qfunc_output)(r) for r in res]
                     res = tuple(res)
diff --git a/pennylane/return_types.py b/pennylane/return_types.py
index 7f9c68014..161615cb5 100644
--- a/pennylane/return_types.py
+++ b/pennylane/return_types.py
@@ -151,7 +151,7 @@ def enable_return():
 
             tape.jacobian(res, [a, b])
 
-        If the measurements do not have the same shape then you need to use concatenation:
+        If the measurements do not have the same shape then you need to use TF hstack:
 
         .. code-block:: python
 
@@ -169,7 +169,7 @@ def enable_return():
 
             with tf.GradientTape() as tape:
                 res = circuit(a, b)
-                res = tf.concat([tf.reshape(i, [-1]) for i in res], 0)
+                res = tf.experimental.numpy.hstack(res)
 
             assert circuit.qtape.trainable_params == [0, 1]
 


Test patch:
diff --git a/tests/returntypes/test_new_return_types_qnode.py b/tests/returntypes/test_new_return_types_qnode.py
index 4c03b0a52..6b4dbe638 100644
--- a/tests/returntypes/test_new_return_types_qnode.py
+++ b/tests/returntypes/test_new_return_types_qnode.py
@@ -2660,7 +2660,11 @@ class TestIntegrationJacobianBackpropMultipleReturns:
             return qml.expval(qml.Projector([0], wires=0)), qml.expval(qml.PauliZ(wires=1))
 
         x = qml.numpy.array([0.1, 0.2, 0.3], requires_grad=True)
-        res = qml.jacobian(circuit)(x)
+
+        def cost(a):
+            return qml.numpy.hstack(circuit(a))
+
+        res = qml.jacobian(cost)(x)
 
         assert isinstance(res, np.ndarray)
         assert res.shape == (2, 3)
@@ -2713,6 +2717,7 @@ class TestIntegrationJacobianBackpropMultipleReturns:
 
         with tf.GradientTape() as tape:
             out = circuit(x)
+            out = tf.stack(out)
 
         res = tape.jacobian(out, x)
 
@@ -2742,6 +2747,7 @@ class TestIntegrationJacobianBackpropMultipleReturns:
 
         with tf.GradientTape() as tape:
             out = circuit(x)
+            out = tf.stack(out)
 
         res = tape.jacobian(out, x)
 
@@ -2815,7 +2821,11 @@ class TestIntegrationJacobianBackpropMultipleReturns:
             return qml.probs(op=qml.PauliZ(wires=0)), qml.probs(wires=[1])
 
         x = qml.numpy.array([0.1, 0.2, 0.3], requires_grad=True)
-        res = qml.jacobian(circuit)(x)
+
+        def cost(a):
+            return qml.numpy.stack(circuit(a))
+
+        res = qml.jacobian(cost)(x)
 
         assert isinstance(res, np.ndarray)
         assert res.shape == (2, 2, 3)
@@ -2868,6 +2878,7 @@ class TestIntegrationJacobianBackpropMultipleReturns:
 
         with tf.GradientTape() as tape:
             out = circuit(x)
+            out = tf.stack(out)
 
         res = tape.jacobian(out, x)
 
@@ -2943,7 +2954,11 @@ class TestIntegrationJacobianBackpropMultipleReturns:
             return qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=[0, 1]), qml.vn_entropy(wires=1)
 
         x = qml.numpy.array([0.1, 0.2, 0.3], requires_grad=True)
-        res = qml.jacobian(circuit)(x)
+
+        def cost(a):
+            return qml.numpy.hstack(circuit(a))
+
+        res = qml.jacobian(cost)(x)
 
         assert isinstance(res, np.ndarray)
         assert res.shape == (6, 3)
@@ -3005,6 +3020,7 @@ class TestIntegrationJacobianBackpropMultipleReturns:
 
         with tf.GradientTape() as tape:
             out = circuit(x)
+            out = tf.experimental.numpy.hstack(out)
 
         res = tape.jacobian(out, x)
 
diff --git a/tests/returntypes/test_parameter_shift_hessian_new.py b/tests/returntypes/test_parameter_shift_hessian_new.py
index 4c1f4641e..28e2ebcf4 100644
--- a/tests/returntypes/test_parameter_shift_hessian_new.py
+++ b/tests/returntypes/test_parameter_shift_hessian_new.py
@@ -1002,9 +1002,12 @@ class TestParameterShiftHessianQNode:
             qml.CNOT(wires=[0, 1])
             return qml.probs(wires=0), qml.probs(wires=1)
 
+        def cost(x):
+            return qml.math.stack(circuit(x))
+
         x = np.ones([2], requires_grad=True)
 
-        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        expected = qml.jacobian(qml.jacobian(cost))(x)
         hessian = qml.gradients.param_shift_hessian(circuit)(x)
 
         assert np.allclose(qml.math.transpose(expected, (0, 3, 2, 1)), hessian)
@@ -1024,9 +1027,12 @@ class TestParameterShiftHessianQNode:
             qml.RY(x[0, 0], wires=0)
             return qml.probs(wires=0), qml.probs(wires=1)
 
+        def cost(x):
+            return qml.math.stack(circuit(x))
+
         x = np.ones([1, 3], requires_grad=True)
 
-        expected = qml.jacobian(qml.jacobian(circuit))(x)
+        expected = qml.jacobian(qml.jacobian(cost))(x)
         hessian = qml.gradients.param_shift_hessian(circuit)(x)
 
         assert np.allclose(qml.math.transpose(expected, (0, 2, 3, 4, 5, 1)), hessian)
@@ -1128,12 +1134,15 @@ class TestParameterShiftHessianQNode:
             qml.CRY(y[0, 1], wires=[0, 1])
             return qml.probs(wires=0), qml.probs(wires=1)
 
+        def cost(x, y, z):
+            return qml.math.stack(circuit(x, y, z))
+
         x = np.array(0.1, requires_grad=True)
         y = np.array([[0.5, 0.6], [0.2, 0.1]], requires_grad=True)
         z = np.array([0.3, 0.4], requires_grad=True)
 
         expected = tuple(
-            qml.jacobian(qml.jacobian(circuit, argnum=i), argnum=i)(x, y, z) for i in range(3)
+            qml.jacobian(qml.jacobian(cost, argnum=i), argnum=i)(x, y, z) for i in range(3)
         )
         hessian = qml.gradients.param_shift_hessian(circuit)(x, y, z)
 
diff --git a/tests/returntypes/test_tensorflow_new.py b/tests/returntypes/test_tensorflow_new.py
index 8c0e108ea..cb16158f8 100644
--- a/tests/returntypes/test_tensorflow_new.py
+++ b/tests/returntypes/test_tensorflow_new.py
@@ -673,7 +673,7 @@ class TestTensorFlowExecuteIntegration:
                 qml.probs(wires=[1])
 
             res = execute([tape], dev, **execute_kwargs)[0]
-            res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+            res = tf.experimental.numpy.hstack(res)
 
         expected = np.array(
             [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2]
diff --git a/tests/returntypes/test_tensorflow_qnode_new.py b/tests/returntypes/test_tensorflow_qnode_new.py
index f4f5500fa..2d36f3624 100644
--- a/tests/returntypes/test_tensorflow_qnode_new.py
+++ b/tests/returntypes/test_tensorflow_qnode_new.py
@@ -662,12 +662,11 @@ class TestQubitIntegration:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return [qml.expval(qml.PauliZ(0)), qml.probs(wires=[1])]
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1])
 
         with tf.GradientTape() as tape:
             res = circuit(x, y)
-            if diff_method != "backprop":
-                res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+            res = tf.experimental.numpy.hstack(res)
 
         expected = np.array(
             [
@@ -872,8 +871,7 @@ class TestQubitIntegration:
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
                 res = circuit(x)
-                if diff_method != "backprop":
-                    res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+                res = tf.experimental.numpy.hstack(res)
             g = tape2.jacobian(res, x, experimental_use_pfor=False)
 
         hess = tape1.jacobian(g, x)
@@ -1520,11 +1518,11 @@ class TestAutograph:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return [qml.expval(qml.PauliZ(0)), qml.probs(wires=[1])]
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1])
 
         with tf.GradientTape() as tape:
             res = circuit(x, y)
-            res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+            res = tf.experimental.numpy.hstack(res)
 
         expected = np.array(
             [
@@ -1786,8 +1784,7 @@ class TestReturn:
 
         with tf.GradientTape() as tape:
             res = circuit(a)
-            if diff_method != "backprop":
-                res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+            res = tf.experimental.numpy.hstack(res)
 
         jac = tape.jacobian(res, a)
 
@@ -1813,8 +1810,7 @@ class TestReturn:
 
         with tf.GradientTape() as tape:
             res = circuit(a, b)
-            if diff_method != "backprop":
-                res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+            res = tf.experimental.numpy.hstack(res)
 
         jac = tape.jacobian(res, (a, b))
 
@@ -1845,8 +1841,7 @@ class TestReturn:
 
         with tf.GradientTape() as tape:
             res = circuit(a)
-            if diff_method != "backprop":
-                res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+            res = tf.experimental.numpy.hstack(res)
 
         jac = tape.jacobian(res, a)
 
@@ -1998,8 +1993,7 @@ class TestReturn:
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
                 res = circuit(par_0, par_1)
-                if diff_method != "backprop":
-                    res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+                res = tf.experimental.numpy.hstack(res)
 
             grad = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False)
             grad = tf.concat(grad, 0)
@@ -2035,8 +2029,7 @@ class TestReturn:
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
                 res = circuit(params)
-                if diff_method != "backprop":
-                    res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+                res = tf.experimental.numpy.hstack(res)
 
             grad = tape2.jacobian(res, params, experimental_use_pfor=False)
 
@@ -2065,8 +2058,7 @@ class TestReturn:
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
                 res = circuit(par_0, par_1)
-                if diff_method != "backprop":
-                    res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+                res = tf.experimental.numpy.hstack(res)
 
             grad = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False)
             grad = tf.concat(grad, 0)
@@ -2101,8 +2093,7 @@ class TestReturn:
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
                 res = circuit(params)
-                if diff_method != "backprop":
-                    res = tf.concat([tf.expand_dims(res[0], 0), res[1]], 0)
+                res = tf.experimental.numpy.hstack(res)
 
             grad = tape2.jacobian(res, params, experimental_use_pfor=False)