Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Inconsistent behavior between lighning.qubit and default.qubit
### Expected behavior

The following code
```
dev1 = qml.device("lightning.qubit", wires=1)

@qml.qnode(dev1)
def circuit2(phi1, phi2):
    qml.RX(phi1, wires=0)
    qml.RY(phi2, wires=0)
    return qml.expval(qml.PauliZ(0))

phi1 = np.array([0.54])
phi2 = np.array([0.12])
circuit2(phi1, phi2)
```
should return similar results no matter whether we use `lightning.qubit` or `default.qubit` as the device.

### Actual behavior

However, if the device is `lightning.qubit`, the output is `Array(0.85154057, dtype=float32, weak_type=True)` while if we use `default.qubit` the result is `Array([0.85154057], dtype=float32)`.

### Additional information

_No response_

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.37.0.dev0
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/user/miniforge3/envs/vw/lib/python3.12/site-packages
Editable project location: /Users/pablo.casares/Developer/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-qiskit, PennyLane_Lightning

Platform info:           macOS-14.5-arm64-arm-64bit
Python version:          3.12.3
Numpy version:           1.26.4
Scipy version:           1.13.0
Installed devices:
- default.clifford (PennyLane-0.37.0.dev0)
- default.gaussian (PennyLane-0.37.0.dev0)
- default.mixed (PennyLane-0.37.0.dev0)
- default.qubit (PennyLane-0.37.0.dev0)
- default.qubit.autograd (PennyLane-0.37.0.dev0)
- default.qubit.jax (PennyLane-0.37.0.dev0)
- default.qubit.legacy (PennyLane-0.37.0.dev0)
- default.qubit.tf (PennyLane-0.37.0.dev0)
- default.qubit.torch (PennyLane-0.37.0.dev0)
- default.qutrit (PennyLane-0.37.0.dev0)
- default.qutrit.mixed (PennyLane-0.37.0.dev0)
- default.tensor (PennyLane-0.37.0.dev0)
- null.qubit (PennyLane-0.37.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.36.0)
- qiskit.basicaer (PennyLane-qiskit-0.36.0)
- qiskit.basicsim (PennyLane-qiskit-0.36.0)
- qiskit.ibmq (PennyLane-qiskit-0.36.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.36.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.36.0)
- qiskit.remote (PennyLane-qiskit-0.36.0)
- lightning.qubit (PennyLane_Lightning-0.37.0.dev49)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request:
Treat batch size of 1 as a batch size
**Context:**

Our result type spec does often squeeze out singleton dimensions, such as those arising from a single measurement, or a shot vector of length 1. 

But we have an inconsistency around whether or not a batch size of `1`, `op.batch_size == 1` is considered an actual batch size or a dimension to squeeze out. I would consider a batch size of 1 different than a batch_size of `None`, and worth maintaining.  This was already the behaviour of `default.qubit`, and is generally how parameter broadcasting works.

**Description of the Change:**

Eliminates squeezing steps from `broadcast_expand`.

**Benefits:**

Consistent, predictable handling of batch dimensions.

**Possible Drawbacks:**

Someone might have been depending on the squeezing behaviour. That code might now break. This may be risky, especially this close to the release.

**Related GitHub Issues:**

[sc-68244] Fixes #5977 


Patch:
diff --git a/doc/releases/changelog-0.38.0.md b/doc/releases/changelog-0.38.0.md
index a96ce8252..5e1f1ec09 100644
--- a/doc/releases/changelog-0.38.0.md
+++ b/doc/releases/changelog-0.38.0.md
@@ -375,6 +375,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.transforms.broadcast_expand` no longer squeezes out batch sizes of size 1, as a batch size of 1 is still a
+  batch size.
+  [(#6147)](https://github.com/PennyLaneAI/pennylane/pull/6147)
+
 * Catalyst replaced `argnum` with `argnums` in gradient related functions, therefore we update the Catalyst
   calls to those functions in PennyLane.
   [(#6117)](https://github.com/PennyLaneAI/pennylane/pull/6117)
diff --git a/pennylane/transforms/broadcast_expand.py b/pennylane/transforms/broadcast_expand.py
index 9ce8caef9..0292af60d 100644
--- a/pennylane/transforms/broadcast_expand.py
+++ b/pennylane/transforms/broadcast_expand.py
@@ -48,6 +48,13 @@ def _split_operations(ops, num_tapes):
     return new_ops
 
 
+def null_postprocessing(results):
+    """A postprocesing function returned by a transform that only converts the batch of results
+    into a result for a single ``QuantumTape``.
+    """
+    return results[0]
+
+
 @transform
 def broadcast_expand(tape: QuantumScript) -> tuple[QuantumScriptBatch, PostprocessingFn]:
     r"""Expand a broadcasted tape into multiple tapes
@@ -124,69 +131,56 @@ def broadcast_expand(tape: QuantumScript) -> tuple[QuantumScriptBatch, Postproce
     tensor([0.98006658, 0.82533561, 0.54030231], requires_grad=True)
     """
     if tape.batch_size is None:
-        output_tapes = [tape]
-
-        def null_postprocessing(results):
-            """A postprocesing function returned by a transform that only converts the batch of results
-            into a result for a single ``QuantumTape``.
-            """
-            return results[0]
+        return (tape,), null_postprocessing
+
+    has_postselect = any(
+        op.postselect is not None for op in tape.operations if isinstance(op, MidMeasureMP)
+    )
+    has_sample = any(isinstance(op, SampleMP) for op in tape.measurements)
+    if has_postselect and has_sample:
+        raise ValueError(
+            "Returning qml.sample is not supported when using post-selected mid-circuit measurements and parameters broadcasting."
+        )
 
-        processing_fn = null_postprocessing
-    else:
+    num_tapes = tape.batch_size
+    new_ops = _split_operations(tape.operations, num_tapes)
 
-        has_postselect = any(
-            op.postselect is not None for op in tape.operations if isinstance(op, MidMeasureMP)
+    output_tapes = tuple(
+        qml.tape.QuantumScript(
+            ops, tape.measurements, shots=tape.shots, trainable_params=tape.trainable_params
         )
-        has_sample = any(isinstance(op, SampleMP) for op in tape.measurements)
-        if has_postselect and has_sample:
-            raise ValueError(
-                "Returning qml.sample is not supported when using post-selected mid-circuit measurements and parameters broadcasting."
-            )
+        for ops in new_ops
+    )
 
-        num_tapes = tape.batch_size
-        new_ops = _split_operations(tape.operations, num_tapes)
+    def processing_fn(results: qml.typing.ResultBatch) -> qml.typing.Result:
+        # closure variables: tape.shots, tape.batch_size, tape.measurements
 
-        output_tapes = []
-        for ops in new_ops:
-            new_tape = qml.tape.QuantumScript(
-                ops, tape.measurements, shots=tape.shots, trainable_params=tape.trainable_params
-            )
-            output_tapes.append(new_tape)
-
-        def processing_fn(results: qml.typing.ResultBatch) -> qml.typing.Result:
-            # The shape of the results should be as follows: results[s][m][b], where s is the shot
-            # vector index, m is the measurement index, and b is the batch index. The shape that
-            # the processing function receives is results[b][s][m].
-
-            if tape.shots.has_partitioned_shots:
-                if len(tape.measurements) > 1:
-                    return tuple(
-                        tuple(
-                            qml.math.squeeze(
-                                qml.math.stack([results[b][s][m] for b in range(tape.batch_size)])
-                            )
-                            for m in range(len(tape.measurements))
-                        )
-                        for s in range(tape.shots.num_copies)
-                    )
+        # The shape of the results should be as follows: results[s][m][b], where s is the shot
+        # vector index, m is the measurement index, and b is the batch index. The shape that
+        # the processing function receives is results[b][s][m].
 
-                # Only need to transpose results[b][s] -> results[s][b]
+        if tape.shots.has_partitioned_shots:
+            if len(tape.measurements) > 1:
                 return tuple(
-                    qml.math.squeeze(
-                        qml.math.stack([results[b][s] for b in range(tape.batch_size)])
+                    tuple(
+                        qml.math.stack([results[b][s][m] for b in range(tape.batch_size)])
+                        for m in range(len(tape.measurements))
                     )
                     for s in range(tape.shots.num_copies)
                 )
 
-            if len(tape.measurements) > 1:
-                # Only need to transpose results[b][m] -> results[m][b]
-                return tuple(
-                    qml.math.squeeze(
-                        qml.math.stack([results[b][m] for b in range(tape.batch_size)])
-                    )
-                    for m in range(len(tape.measurements))
-                )
-            return qml.math.squeeze(qml.math.stack(results))
+            # Only need to transpose results[b][s] -> results[s][b]
+            return tuple(
+                qml.math.stack([results[b][s] for b in range(tape.batch_size)])
+                for s in range(tape.shots.num_copies)
+            )
+
+        if len(tape.measurements) > 1:
+            # Only need to transpose results[b][m] -> results[m][b]
+            return tuple(
+                qml.math.stack([results[b][m] for b in range(tape.batch_size)])
+                for m in range(len(tape.measurements))
+            )
+        return qml.math.stack(results)
 
     return output_tapes, processing_fn
diff --git a/pennylane/workflow/return_types_spec.rst b/pennylane/workflow/return_types_spec.rst
index e5e5a69a5..6e9ba4c91 100644
--- a/pennylane/workflow/return_types_spec.rst
+++ b/pennylane/workflow/return_types_spec.rst
@@ -71,7 +71,8 @@ Broadcasting
 Parameter broadcasting adds a leading dimension to the numeric array itself.
 
 If the corresponding tape has a ``batch_size`` and the result object is numeric, then the numeric object should
-gain a leading dimension.
+gain a leading dimension.  Note that a batch size of ``1`` is still a batch size,
+and still should correspond to a leading dimension.
 
 >>> op = qml.RX((0, np.pi/4, np.pi/2), wires=0)
 >>> tape = qml.tape.QuantumScript((op,), [qml.probs(wires=0)])


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit_preprocessing.py b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
index cb1e87070..29a1fdca7 100644
--- a/tests/devices/default_qubit/test_default_qubit_preprocessing.py
+++ b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
@@ -437,7 +437,7 @@ class TestPreprocessingIntegration:
     def test_preprocess_batch_transform_adjoint(self):
         """Test that preprocess returns the correct tapes when a batch transform
         is needed."""
-        ops = [qml.Hadamard(0), qml.CNOT([0, 1]), qml.RX([np.pi, np.pi / 2], wires=1)]
+        ops = [qml.Hadamard(0), qml.CNOT([0, 1]), qml.RX([np.pi, np.pi / 2, 2.5], wires=1)]
         # Need to specify grouping type to transform tape
         measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))]
         tapes = [
@@ -453,20 +453,23 @@ class TestPreprocessingIntegration:
         expected_ops = [
             [qml.Hadamard(0), qml.CNOT([0, 1]), qml.RX(np.pi, wires=1)],
             [qml.Hadamard(0), qml.CNOT([0, 1]), qml.RX(np.pi / 2, wires=1)],
+            [qml.Hadamard(0), qml.CNOT([0, 1]), qml.RX(2.5, wires=1)],
         ]
 
-        assert len(res_tapes) == 4
+        assert len(res_tapes) == 6
         for i, t in enumerate(res_tapes):
-            for op, expected_op in zip(t.operations, expected_ops[i % 2]):
+            for op, expected_op in zip(t.operations, expected_ops[i % 3]):
                 qml.assert_equal(op, expected_op)
             assert len(t.measurements) == 1
-            if i < 2:
+            if i < 3:
                 qml.assert_equal(t.measurements[0], measurements[0])
             else:
                 qml.assert_equal(t.measurements[0], measurements[1])
 
-        val = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]])
-        assert np.array_equal(batch_fn(val), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]))
+        # outer dimension = tapes, each has one meausrement
+        val = (1, 2, 3, 4, 5, 6)
+        expected = (np.array([1, 2, 3]), np.array([4, 5, 6]))
+        assert np.array_equal(batch_fn(val), expected)
 
     def test_preprocess_expand(self):
         """Test that preprocess returns the correct tapes when expansion is needed."""
@@ -526,7 +529,7 @@ class TestPreprocessingIntegration:
     def test_preprocess_split_and_expand_adjoint(self):
         """Test that preprocess returns the correct tapes when splitting and expanding
         is needed."""
-        ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2], wires=1)]
+        ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2, 1.23], wires=1)]
         # Need to specify grouping type to transform tape
         measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))]
         tapes = [
@@ -542,20 +545,22 @@ class TestPreprocessingIntegration:
         expected_ops = [
             [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi, wires=1)],
             [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi / 2, wires=1)],
+            [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(1.23, wires=1)],
         ]
 
-        assert len(res_tapes) == 4
+        assert len(res_tapes) == 6
         for i, t in enumerate(res_tapes):
-            for op, expected_op in zip(t.operations, expected_ops[i % 2]):
+            for op, expected_op in zip(t.operations, expected_ops[i % 3]):
                 qml.assert_equal(op, expected_op)
             assert len(t.measurements) == 1
-            if i < 2:
+            if i < 3:
                 qml.assert_equal(t.measurements[0], measurements[0])
             else:
                 qml.assert_equal(t.measurements[0], measurements[1])
 
-        val = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]])
-        assert np.array_equal(batch_fn(val), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]))
+        val = (1, 2, 3, 4, 5, 6)
+        expected = (np.array([1, 2, 3]), np.array([4, 5, 6]))
+        assert np.array_equal(batch_fn(val), expected)
 
     def test_preprocess_check_validity_fail(self):
         """Test that preprocess throws an error if the split and expanded tapes have
diff --git a/tests/interfaces/test_execute.py b/tests/interfaces/test_execute.py
index 7cfebec32..2fb8206dc 100644
--- a/tests/interfaces/test_execute.py
+++ b/tests/interfaces/test_execute.py
@@ -17,90 +17,9 @@ from contextlib import nullcontext
 import pytest
 
 import pennylane as qml
-from pennylane import numpy as np
 from pennylane.devices import DefaultQubit
 
 
-class TestBatchTransformHelper:
-    """Unit tests for the _batch_transform helper function."""
-
-    def test_warns_if_requested_off(self):
-        """Test that a warning is raised if the the batch transform is requested to not be used."""
-
-        # pylint: disable=too-few-public-methods
-        class CustomOp(qml.operation.Operator):
-            """Dummy operator."""
-
-            def decomposition(self):
-                return [qml.PauliX(self.wires[0])]
-
-        dev = DefaultQubit()
-
-        qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))])
-
-        with pytest.warns(UserWarning, match="Device batch transforms cannot be turned off"):
-            program, _ = dev.preprocess()
-            with pytest.warns(
-                qml.PennyLaneDeprecationWarning,
-                match="The device_batch_transform argument is deprecated",
-            ):
-                qml.execute(
-                    (qs, qs), device=dev, device_batch_transform=False, transform_program=program
-                )
-
-    def test_split_and_expand_performed(self):
-        """Test that preprocess returns the correct tapes when splitting and expanding
-        is needed."""
-
-        class NoMatOp(qml.operation.Operation):
-            """Dummy operation for expanding circuit."""
-
-            # pylint: disable=missing-function-docstring
-            num_wires = 1
-
-            # pylint: disable=arguments-renamed, invalid-overridden-method
-            @property
-            def has_matrix(self):
-                return False
-
-            def decomposition(self):
-                return [qml.PauliX(self.wires), qml.PauliY(self.wires)]
-
-        ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2], wires=1)]
-        # Need to specify grouping type to transform tape
-        measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))]
-        tapes = [
-            qml.tape.QuantumScript(ops=ops, measurements=[measurements[0]]),
-            qml.tape.QuantumScript(ops=ops, measurements=[measurements[1]]),
-        ]
-
-        dev = DefaultQubit()
-        config = qml.devices.ExecutionConfig(gradient_method="adjoint")
-
-        program, new_config = dev.preprocess(config)
-        res_tapes, batch_fn = program(tapes)
-        expected_ops = [
-            [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi, wires=1)],
-            [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi / 2, wires=1)],
-        ]
-
-        assert len(res_tapes) == 4
-        for i, t in enumerate(res_tapes):
-            for op, expected_op in zip(t.operations, expected_ops[i % 2]):
-                qml.assert_equal(op, expected_op)
-            assert len(t.measurements) == 1
-            if i < 2:
-                qml.assert_equal(t.measurements[0], measurements[0])
-            else:
-                qml.assert_equal(t.measurements[0], measurements[1])
-
-        input = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]])
-        assert np.array_equal(batch_fn(input), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]))
-
-        assert new_config.grad_on_execution
-        assert new_config.use_device_gradient
-
-
 def test_warning_if_not_device_batch_transform():
     """Test that a warning is raised if the users requests to not run device batch transform."""
 
diff --git a/tests/transforms/test_broadcast_expand.py b/tests/transforms/test_broadcast_expand.py
index 935e8cb56..e1212626f 100644
--- a/tests/transforms/test_broadcast_expand.py
+++ b/tests/transforms/test_broadcast_expand.py
@@ -56,13 +56,11 @@ H0 = qml.Hamiltonian(qml.math.array(coeffs0), [qml.PauliZ(0), qml.PauliY(1)])
 
 # Here we exploit the product structure of our circuit
 def exp_fn_Z0(x, y, z):
-    out = -qml.math.cos(x) * qml.math.ones_like(y) * qml.math.ones_like(z)
-    return out[0] if len(out) == 1 else out
+    return -qml.math.cos(x) * qml.math.ones_like(y) * qml.math.ones_like(z)
 
 
 def exp_fn_Y1(x, y, z):
-    out = qml.math.sin(y) * qml.math.cos(z) * qml.math.ones_like(x)
-    return out[0] if len(out) == 1 else out
+    return qml.math.sin(y) * qml.math.cos(z) * qml.math.ones_like(x)
 
 
 def exp_fn_Z0Y1(x, y, z):