Shot batching is made more efficient by executing all the shots in one go on Lightning Qubit

.github/CHANGELOG.md

@@ -20,10 +20,10 @@
 
 * Avoid unnecessary memory reset in LightningQubit's state vector class constructor.
   [(#811)](https://github.com/PennyLaneAI/pennylane-lightning/pull/811)
-  
+
 * Add `generate_samples(wires)` support in LightningQubit.
   [(#809)](https://github.com/PennyLaneAI/pennylane-lightning/pull/809)
-  
+
 * Optimize the OpenMP parallelization of Lightning-Qubit's `probs` for all number of targets.
   [(#807)](https://github.com/PennyLaneAI/pennylane-lightning/pull/807)
 
@@ -48,6 +48,9 @@
 * Add a Catalyst-specific wrapping class for Lightning Kokkos.
   [(#770)](https://github.com/PennyLaneAI/pennylane-lightning/pull/770)
 
+* Shot batching is made more efficient by executing all the shots in one go on LightningQubit.
+  [#814](https://github.com/PennyLaneAI/pennylane-lightning/pull/814)
+
 ### Documentation
 
 ### Bug fixes
@@ -71,7 +74,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Amintor Dusko, Vincent Michaud-Rioux, Shuli Shu
+Ali Asadi, Amintor Dusko, Vincent Michaud-Rioux, Shuli Shu
 
 ---
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.38.0-dev16"
+__version__ = "0.38.0-dev17"

pennylane_lightning/lightning_qubit/_measurements.py

@@ -401,31 +401,6 @@
 
             return tuple(processed)
 
-        # if there is a shot vector, build a list containing results for each shot entry
-        if shots.has_partitioned_shots:
-            processed_samples = []
-            for s in shots:
-                # currently we call sample_state for each shot entry, but it may be
-                # better to call sample_state just once with total_shots, then use
-                # the shot_range keyword argument
-                try:
-                    if self._mcmc:
-                        samples = self._measurement_lightning.generate_mcmc_samples(
-                            len(wires), self._kernel_name, self._num_burnin, s
-                        ).astype(int, copy=False)
-                    else:
-                        samples = self._measurement_lightning.generate_samples(
-                            list(wires), s
-                        ).astype(int, copy=False)
-                except ValueError as e:
-                    if str(e) != "probabilities contain NaN":
-                        raise e
-                    samples = qml.math.full((s, len(wires)), 0)
-
-                processed_samples.append(_process_single_shot(samples))
-            self._apply_diagonalizing_gates(mps, adjoint=True)
-            return tuple(zip(*processed_samples))
-
         try:
             if self._mcmc:
                 samples = self._measurement_lightning.generate_mcmc_samples(
@@ -442,7 +417,16 @@
 
         self._apply_diagonalizing_gates(mps, adjoint=True)
 
-        return _process_single_shot(samples)
+        # if there is a shot vector, use the shots.bins generator to
+        # split samples w.r.t. the shots
+        processed_samples = []
+        for lower, upper in shots.bins():
+            shot = _process_single_shot(samples[..., lower:upper, :])
+            processed_samples.append(shot)
+
+        return (
+            tuple(zip(*processed_samples)) if shots.has_partitioned_shots else processed_samples[0]
+        )
 
     def _measure_hamiltonian_with_samples(
         self,

tests/test_measurements.py

@@ -23,7 +23,7 @@
 from conftest import LightningDevice as ld
 from conftest import device_name, lightning_ops, validate_measurements
 from flaky import flaky
-from pennylane.measurements import Expectation, Variance
+from pennylane.measurements import Expectation, Shots, Variance
 
 if not ld._CPP_BINARY_AVAILABLE:
     pytest.skip("No binary module found. Skipping.", allow_module_level=True)
@@ -828,3 +828,21 @@ def func(x, y):
     results2 = func2(*params)
 
     validate_measurements(measure_f, shots, results1, results2)
+
+
+# TODO: Add LK and LGPU after migrating to the new device API
+# TODO: Add LT after extending the support for shots_vector
+@pytest.mark.skipif(
+    device_name != "lightning.qubit",
+    reason="lightning.qubit only has support for vector_shots",
+)
+@pytest.mark.parametrize("shots", ((1, 10), (1, 10, 100), (1, 10, 10, 100, 100, 100)))
+def test_shots_bins(shots, qubit_device):
+    """Tests that Lightning handles multiple shots."""
+
+    @qml.qnode(qubit_device(wires=1, shots=shots))
+    def circuit():
+        return qml.expval(qml.PauliZ(wires=0))
+
+    assert np.sum(shots) == circuit.device.shots.total_shots
+    assert np.allclose(circuit(), 1.0)