From 559b1bf47922e7113744373650720bac2a5e3841 Mon Sep 17 00:00:00 2001
From: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>
Date: Fri, 25 Jun 2021 13:18:37 -0400
Subject: [PATCH] Fix return of single measurement in qfunc transform. (#1434)

* Fix return of single measurement in qfunc transform.

* Add unit test.

* Update CHANGELOG.

Co-authored-by: Olivia Di Matteo <dimatteo.olivia@gmail.com>
---
 .github/CHANGELOG.md                     |  6 ++++-
 pennylane/transforms/qfunc_transforms.py |  4 +++
 tests/transforms/test_qfunc_transform.py | 31 ++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 30b4da82720..8bbec02dd39 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -11,13 +11,17 @@
 
 <h3>Bug fixes</h3>
 
+* Quantum function transforms now preserve the format of the measurement
+  results, so that a single measurement returns a single value rather than
+  an array with a single element. [(#1434)](https://github.com/PennyLaneAI/pennylane/pull/1434/files)
+
 <h3>Documentation</h3>
 
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
 
-Ashish Panigrahi
+Olivia Di Matteo, Ashish Panigrahi
 
 
 # Release 0.16.0 (current release)
diff --git a/pennylane/transforms/qfunc_transforms.py b/pennylane/transforms/qfunc_transforms.py
index 9581bae67aa..13263af3075 100644
--- a/pennylane/transforms/qfunc_transforms.py
+++ b/pennylane/transforms/qfunc_transforms.py
@@ -182,6 +182,10 @@ def _create_qfunc_internal_wrapper(fn, tape_transform, transform_args, transform
     def internal_wrapper(*args, **kwargs):
         tape = make_tape(fn)(*args, **kwargs)
         tape = tape_transform(tape, *transform_args, **transform_kwargs)
+
+        if len(tape.measurements) == 1:
+            return tape.measurements[0]
+
         return tape.measurements
 
     return internal_wrapper
diff --git a/tests/transforms/test_qfunc_transform.py b/tests/transforms/test_qfunc_transform.py
index f7d157de5ae..066a24f3889 100644
--- a/tests/transforms/test_qfunc_transform.py
+++ b/tests/transforms/test_qfunc_transform.py
@@ -271,6 +271,37 @@ def ansatz():
         assert ops[0].parameters == [x]
         assert ops[1].name == "CZ"
 
+    def test_transform_single_measurement(self):
+        """Test that transformed functions return a scalar value when there is only
+        a single measurement."""
+
+        @qml.qfunc_transform
+        def expand_hadamards(tape):
+            for op in tape.operations + tape.measurements:
+                if op.name == "Hadamard":
+                    qml.RZ(np.pi, wires=op.wires)
+                    qml.RY(np.pi / 2, wires=op.wires)
+                else:
+                    op.queue()
+
+        def ansatz():
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(wires=1))
+
+        dev = qml.device("default.qubit", wires=2)
+
+        normal_qnode = qml.QNode(ansatz, dev)
+
+        transformed_ansatz = expand_hadamards(ansatz)
+        transformed_qnode = qml.QNode(transformed_ansatz, dev)
+
+        normal_result = normal_qnode()
+        transformed_result = transformed_qnode()
+
+        assert np.allclose(normal_result, transformed_result)
+        assert normal_result.shape == transformed_result.shape
+
 
 ############################################
 # Test transform, ansatz, and qfunc function