Include the changes introduced in #814 for measurements

pennylane_lightning/lightning_kokkos/_measurements.py

@@ -388,39 +388,6 @@ def _process_single_shot(samples):
 
             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:
-                    #  ----------------------------------------
-                    # Original:
-                    # 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)
-
-                    # Specific for Kokkos:
-                    samples = self._measurement_lightning.generate_samples(len(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:
             #  ----------------------------------------
             # Original:
@@ -446,7 +413,16 @@ def _process_single_shot(samples):
 
         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():
+            result = _process_single_shot(samples[..., lower:upper, :])
+            processed_samples.append(result)
+
+        return (
+            tuple(zip(*processed_samples)) if shots.has_partitioned_shots else processed_samples[0]
+        )
 
     def _measure_hamiltonian_with_samples(
         self,