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

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

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

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