No longer squeeze singleton dimension in samples (#7944)

**Context:**

Similar to #7917 , but doesn't squeeze observables.

Currently, we always squeeze out singleton dimension from `qml.sample`
measurments. This leads to lots of additional logical branches, and
error-prone postprocessing. It can be really easy to loose track of
which dimensions are being sliced into.

Catalyst also never squeezes out singleton dimensions. If we want to
unify the behavior of pennylane and catalyst, we need to either start
squeezing in catalyst, or stop squeezing in pennylane. We are opting to
stop squeezing in pennylane.

**Description of the Change:**

Switches to that sampling wires *always* returns an array of shape
`(shots, wires)`, regardless of whether or not we have a single shot or
a single wire.

**Benefits:**

Much simpler, more robust code and unified behavior between pennylane
and catalyst.

**Possible Drawbacks:**

This is going to be a breaking change, and there's really no way to
soften this to a deprecation. We just have to pull the bandaid off.

Need an accompanying issue to lightning.

**Related GitHub Issues:**

[[sc-95797](https://app.shortcut.com/xanaduai/story/95797)]

---------

Co-authored-by: Andrija Paurevic <[email protected]>
Co-authored-by: Yushao Chen (Jerry) <[email protected]>
Co-authored-by: Isaac De Vlugt <[email protected]>

Author: 4 people

doc/development/plugins.rst

@@ -87,7 +87,21 @@ circuits.
 >>> tape1 = qml.tape.QuantumScript([], [qml.sample(wires=0)], shots=10)
 >>> dev = qml.device('default.qubit')
 >>> dev.execute((tape0, tape1))
-(array([0, 0, 0, 0, 0]), array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]))
+(array([[0],
+        [0],
+        [0],
+        [0],
+        [0]]),
+ array([[0],
+        [0],
+        [0],
+        [0],
+        [0],
+        [0],
+        [0],
+        [0],
+        [0],
+        [0]]))
 
 The :class:`~.measurements.Shots` class describes the shots. Users can optionally specify a shot vector, or
 different numbers of shots to use when calculating the final expectation value.
@@ -457,9 +471,27 @@ the top user level, we aim to allow dynamic configuration of the device.
 >>> config = qml.devices.ExecutionConfig(device_options={"rng": 42})
 >>> tape = qml.tape.QuantumTape([qml.Hadamard(0)], [qml.sample(wires=0)], shots=10)
 >>> dev.execute(tape, config)
-array([1, 0, 1, 1, 0, 1, 1, 1, 0, 0])
+array([[1],
+       [1],
+       [0],
+       [1],
+       [0],
+       [1],
+       [0],
+       [1],
+       [0],
+       [0]])
 >>> dev.execute(tape, config)
-array([1, 0, 1, 1, 0, 1, 1, 1, 0, 0])
+array([[0],
+       [1],
+       [0],
+       [0],
+       [0],
+       [1],
+       [1],
+       [0],
+       [0],
+       [0]])
 
 By pulling options from this dictionary instead of from device properties, we unlock two key
 pieces of functionality:

doc/introduction/dynamic_quantum_circuits.rst

@@ -94,9 +94,13 @@ Executing this QNode with 10 shots yields
 .. code-block:: pycon
 
     >>> func(np.pi / 2, shots=10)
-    array([1, 1, 1, 1, 1, 1, 1])
+    array([[1],
+       [1],
+       [1],
+       [1],
+       [1]])
 
-Note that only 7 samples are returned. This is because samples that do not meet the postselection criteria are
+Note that less than 10 samples are returned. This is because samples that do not meet the postselection criteria are
 discarded. This behaviour can be customized, see the section
 :ref:`"Configuring mid-circuit measurements" <mcm_config>`.
 

doc/introduction/unsupported_gradients.rst

@@ -287,7 +287,7 @@ error, but the results will be incorrect:
     def sample_backward():
         dev = qml.device('default.qubit', wires=1)
 
-        @partial(qml.set_shots, shots=20)
+        @partial(qml.set_shots, shots=5)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x[0], wires=0)
@@ -297,26 +297,15 @@ error, but the results will be incorrect:
         print(qml.jacobian(circuit)(x))
 
 >>> sample_backward()
-[[0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]
- [0.5]]
+[[[0.5]]
+<BLANKLINE>
+ [[0.5]]
+<BLANKLINE>
+ [[0.5]]
+<BLANKLINE>
+ [[0.5]]
+<BLANKLINE>
+ [[0.5]]]
 
 The forward pass is supported and will work as expected:
 
@@ -335,4 +324,23 @@ The forward pass is supported and will work as expected:
         print(circuit(x))
 
 >>> sample_forward()
-[0 1 0 0 0 1 1 0 0 1 1 1 0 0 0 1 1 0 0 0]
+[[0]
+ [0]
+ [0]
+ [0]
+ [1]
+ [1]
+ [0]
+ [0]
+ [1]
+ [1]
+ [1]
+ [1]
+ [0]
+ [1]
+ [1]
+ [0]
+ [1]
+ [0]
+ [0]
+ [1]]

doc/releases/changelog-dev.md

@@ -291,6 +291,40 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `qml.sample` no longer has singleton dimensions squeezed out for single shots or single wires. This cuts
+  down on the complexity of post-processing due to having to handle single shot and single wire cases
+  separately. The return shape will now *always* be `(shots, num_wires)`.
+  [(#7944)](https://github.com/PennyLaneAI/pennylane/pull/7944)
+
+  For a simple qnode:
+
+  ```pycon
+  >>> @qml.qnode(qml.device('default.qubit'))
+  ... def c():
+  ...   return qml.sample(wires=0)
+  ```
+
+  Before the change, we had:
+  
+  ```pycon
+  >>> qml.set_shots(c, shots=1)()
+  0
+  ```
+
+  and now we have:
+
+  ```pycon
+  >>> qml.set_shots(c, shots=1)()
+  array([[0]])
+  ```
+
+  Previous behavior can be recovered by squeezing the output:
+
+  ```pycon
+  >>> qml.math.squeeze(qml.set_shots(c, shots=1)())
+  0
+  ```
+
 * `ExecutionConfig` and `MCMConfig` from `pennylane.devices` are now frozen dataclasses whose fields should be updated with `dataclass.replace`. 
   [(#7697)](https://github.com/PennyLaneAI/pennylane/pull/7697)
 

pennylane/devices/_qubit_device.py

@@ -659,7 +659,7 @@ def statistics(
             elif isinstance(m, SampleMP):
                 samples = self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=False)
                 dtype = int if isinstance(obs, SampleMP) else None
-                result = self._asarray(qml.math.squeeze(samples), dtype=dtype)
+                result = self._asarray(samples, dtype=dtype)
 
             elif isinstance(m, CountsMP):
                 result = self.sample(m, shot_range=shot_range, bin_size=bin_size, counts=True)

pennylane/devices/legacy_facade.py

@@ -152,7 +152,11 @@ class LegacyDeviceFacade(Device):
     Shots(total_shots=None, shot_vector=())
     >>> tape = qml.tape.QuantumScript([], [qml.sample(wires=0)], shots=5)
     >>> new_dev.execute(tape)
-    array([0, 0, 0, 0, 0])
+    array([[0],
+       [0],
+       [0],
+       [0],
+       [0]])
 
     """
 

pennylane/devices/preprocess.py

@@ -639,31 +639,19 @@ def measurements_from_samples(tape):
     diagonalized_tape, measured_wires = _get_diagonalized_tape_and_wires(tape)
     new_tape = diagonalized_tape.copy(measurements=[qml.sample(wires=measured_wires)])
 
-    def unsqueezed(samples):
-        """If the samples have been squeezed to remove the 'extra' dimension in the case where
-        shots=1 or wires=1, unsqueeze to restore the raw samples format expected by mp.process_samples
-        """
-
-        # if we stop squeezing out the extra dimension for shots=1 or wires=1 when sampling (as is
-        # already the case in Catalyst), this problem goes away
-
-        if len(samples.shape) == 1:
-            samples = qml.math.array([[s] for s in samples], like=samples)
-        return samples
-
     def postprocessing_fn(results):
         """A processing function to get measurement values from samples."""
         samples = results[0]
         if tape.shots.has_partitioned_shots:
             results_processed = []
             for s in samples:
-                res = [m.process_samples(unsqueezed(s), measured_wires) for m in tape.measurements]
+                res = [m.process_samples(s, measured_wires) for m in tape.measurements]
                 if len(tape.measurements) == 1:
                     res = res[0]
                 results_processed.append(res)
         else:
             results_processed = [
-                m.process_samples(unsqueezed(samples), measured_wires) for m in tape.measurements
+                m.process_samples(samples, measured_wires) for m in tape.measurements
             ]
             if len(tape.measurements) == 1:
                 results_processed = results_processed[0]

pennylane/devices/qubit/sampling.py

@@ -19,7 +19,6 @@
 import pennylane as qml
 from pennylane.measurements import (
     ClassicalShadowMP,
-    CountsMP,
     ExpectationMP,
     SampleMeasurement,
     ShadowExpvalMP,
@@ -308,15 +307,7 @@ def _measure_with_samples_diagonalizing_gates(
     wires = qml.wires.Wires(range(total_indices))
 
     def _process_single_shot(samples):
-        processed = []
-        for mp in mps:
-            res = mp.process_samples(samples, wires)
-            if not isinstance(mp, CountsMP):
-                res = qml.math.squeeze(res)
-
-            processed.append(res)
-
-        return tuple(processed)
+        return tuple(mp.process_samples(samples, wires) for mp in mps)
 
     try:
         prng_key, _ = jax_random_split(prng_key)

pennylane/devices/qubit/simulate.py

@@ -858,8 +858,6 @@ def combine_measurements(terminal_measurements, results, mcm_samples):
             comb_meas = measurement_with_no_shots(circ_meas)
         else:
             comb_meas = combine_measurements_core(circ_meas, results.pop(0))
-        if isinstance(circ_meas, SampleMP):
-            comb_meas = qml.math.squeeze(comb_meas)
         final_measurements.append(comb_meas)
     return final_measurements[0] if len(final_measurements) == 1 else tuple(final_measurements)
 
@@ -916,7 +914,7 @@ def _(original_measurement: SampleMP, measures):
     new_sample = tuple(
         qml.math.atleast_1d(m[1]) for m in measures.values() if m[0] and not m[1] is tuple()
     )
-    return qml.math.squeeze(qml.math.concatenate(new_sample))
+    return qml.math.concatenate(new_sample)
 
 
 @debug_logger

pennylane/devices/qubit_mixed/sampling.py

@@ -20,9 +20,8 @@
 import numpy as np
 
 import pennylane as qml
-from pennylane import math
 from pennylane.devices.qubit.sampling import _group_measurements, jax_random_split, sample_probs
-from pennylane.measurements import CountsMP, ExpectationMP, SampleMeasurement, Shots
+from pennylane.measurements import ExpectationMP, SampleMeasurement, Shots
 from pennylane.measurements.classical_shadow import ClassicalShadowMP, ShadowExpvalMP
 from pennylane.ops import LinearCombination, Sum
 from pennylane.typing import TensorLike
@@ -87,15 +86,7 @@ def _measure_with_samples_diagonalizing_gates(
     wires = qml.wires.Wires(range(total_indices))
 
     def _process_single_shot(samples):
-        processed = []
-        for mp in mps:
-            res = mp.process_samples(samples, wires)
-            if not isinstance(mp, CountsMP):
-                res = math.squeeze(res)
-
-            processed.append(res)
-
-        return tuple(processed)
+        return tuple(mp.process_samples(samples, wires) for mp in mps)
 
     prng_key, _ = jax_random_split(prng_key)
     samples = sample_state(