From f577c4e76553723f8a5a23311b9f2c56e9df9873 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Thu, 25 Jul 2024 15:31:35 -0400
Subject: [PATCH 01/26] Added simulate method for the new API with LKokkos

---
 .../lightning_kokkos/_measurements.py         | 501 ++++++++++++++++++
 .../lightning_kokkos/_state_vector.py         | 477 ++++++++++++++++-
 .../lightning_kokkos/lightning_kokkos.py      |  42 +-
 .../test_measurements_class.py                |  24 +-
 .../test_state_vector_class.py                |  15 +-
 5 files changed, 1038 insertions(+), 21 deletions(-)
 create mode 100644 pennylane_lightning/lightning_kokkos/_measurements.py

diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py
new file mode 100644
index 000000000..fc331a799
--- /dev/null
+++ b/pennylane_lightning/lightning_kokkos/_measurements.py
@@ -0,0 +1,501 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Class implementation for state vector measurements.
+"""
+
+# pylint: disable=import-error, no-name-in-module, ungrouped-imports
+try:
+    from pennylane_lightning.lightning_kokkos_ops import MeasurementsC64, MeasurementsC128
+except ImportError:
+    pass
+
+from functools import reduce
+from typing import Callable, List, Union
+
+import numpy as np
+import pennylane as qml
+from pennylane.devices.qubit.sampling import _group_measurements
+from pennylane.measurements import (
+    ClassicalShadowMP,
+    CountsMP,
+    ExpectationMP,
+    MeasurementProcess,
+    ProbabilityMP,
+    SampleMeasurement,
+    ShadowExpvalMP,
+    Shots,
+    StateMeasurement,
+    VarianceMP,
+)
+from pennylane.ops import Hamiltonian, SparseHamiltonian, Sum
+from pennylane.tape import QuantumScript
+from pennylane.typing import Result, TensorLike
+from pennylane.wires import Wires
+
+from pennylane_lightning.core._serialize import QuantumScriptSerializer
+
+
+class LightningKokkosMeasurements:
+    """Lightning Measurements class
+
+    Measures the state provided by the LightningStateVector class.
+
+    Args:
+        qubit_state(LightningStateVector): Lightning state-vector class containing the state vector to be measured.
+        mcmc (bool): Determine whether to use the approximate Markov Chain Monte Carlo
+            sampling method when generating samples.
+        kernel_name (str): name of MCMC transition kernel. The current version supports
+            two kernels: ``"Local"`` and ``"NonZeroRandom"``.
+            The local kernel conducts a bit-flip local transition between states.
+            The local kernel generates a random qubit site and then generates a random
+            number to determine the new bit at that qubit site. The ``"NonZeroRandom"`` kernel
+            randomly transits between states that have nonzero probability.
+        num_burnin (int): number of MCMC steps that will be dropped. Increasing this value will
+            result in a closer approximation but increased runtime.
+    """
+
+    def __init__(
+        self,
+        kokkos_state,
+        mcmc: bool = None,
+        kernel_name: str = None,
+        num_burnin: int = None,
+    ) -> None:
+        self._qubit_state = kokkos_state
+        self._dtype = kokkos_state.dtype
+        self._measurement_lightning = self._measurement_dtype()(kokkos_state.state_vector)
+        self._mcmc = mcmc
+        self._kernel_name = kernel_name
+        self._num_burnin = num_burnin
+        if self._mcmc and not self._kernel_name:
+            self._kernel_name = "Local"
+        if self._mcmc and not self._num_burnin:
+            self._num_burnin = 100
+
+    @property
+    def qubit_state(self):
+        """Returns a handle to the LightningStateVector class."""
+        return self._qubit_state
+
+    @property
+    def dtype(self):
+        """Returns the simulation data type."""
+        return self._dtype
+
+    def _measurement_dtype(self):
+        """Binding to Lightning Measurements C++ class.
+
+        Returns: the Measurements class
+        """
+        return MeasurementsC64 if self.dtype == np.complex64 else MeasurementsC128
+
+    def state_diagonalizing_gates(self, measurementprocess: StateMeasurement) -> TensorLike:
+        """Apply a measurement to state when the measurement process has an observable with diagonalizing gates.
+            This method is bypassing the measurement process to default.qubit implementation.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the state
+
+        Returns:
+            TensorLike: the result of the measurement
+        """
+        diagonalizing_gates = measurementprocess.diagonalizing_gates()
+        self._qubit_state.apply_operations(diagonalizing_gates)
+        state_array = self._qubit_state.state
+        wires = Wires(range(self._qubit_state.num_wires))
+        result = measurementprocess.process_state(state_array, wires)
+        self._qubit_state.apply_operations([qml.adjoint(g) for g in reversed(diagonalizing_gates)])
+        return result
+
+    # pylint: disable=protected-access
+    def expval(self, measurementprocess: MeasurementProcess):
+        """Expectation value of the supplied observable contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the state
+
+        Returns:
+            Expectation value of the observable
+        """
+
+        if isinstance(measurementprocess.obs, qml.SparseHamiltonian):
+            # ensuring CSR sparse representation.
+            CSR_SparseHamiltonian = measurementprocess.obs.sparse_matrix(
+                wire_order=list(range(self._qubit_state.num_wires))
+            ).tocsr(copy=False)
+            return self._measurement_lightning.expval(
+                CSR_SparseHamiltonian.indptr,
+                CSR_SparseHamiltonian.indices,
+                CSR_SparseHamiltonian.data,
+            )
+
+        if (
+            isinstance(measurementprocess.obs, (qml.ops.Hamiltonian, qml.Hermitian))
+            or (measurementprocess.obs.arithmetic_depth > 0)
+            or isinstance(measurementprocess.obs.name, List)
+        ):
+            ob_serialized = QuantumScriptSerializer(
+                self._qubit_state.device_name, self.dtype == np.complex64
+            )._ob(measurementprocess.obs)
+            return self._measurement_lightning.expval(ob_serialized)
+
+        return self._measurement_lightning.expval(
+            measurementprocess.obs.name, measurementprocess.obs.wires
+        )
+
+    def probs(self, measurementprocess: MeasurementProcess):
+        """Probabilities of the supplied observable or wires contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the state
+
+        Returns:
+            Probabilities of the supplied observable or wires
+        """
+        diagonalizing_gates = measurementprocess.diagonalizing_gates()
+        if diagonalizing_gates:
+            self._qubit_state.apply_operations(diagonalizing_gates)
+        results = self._measurement_lightning.probs(measurementprocess.wires.tolist())
+        if diagonalizing_gates:
+            self._qubit_state.apply_operations(
+                [qml.adjoint(g, lazy=False) for g in reversed(diagonalizing_gates)]
+            )
+        return results
+
+    def var(self, measurementprocess: MeasurementProcess):
+        """Variance of the supplied observable contained in the MeasurementProcess.
+
+        Args:
+            measurementprocess (StateMeasurement): measurement to apply to the state
+
+        Returns:
+            Variance of the observable
+        """
+
+        if isinstance(measurementprocess.obs, qml.SparseHamiltonian):
+            # ensuring CSR sparse representation.
+            CSR_SparseHamiltonian = measurementprocess.obs.sparse_matrix(
+                wire_order=list(range(self._qubit_state.num_wires))
+            ).tocsr(copy=False)
+            return self._measurement_lightning.var(
+                CSR_SparseHamiltonian.indptr,
+                CSR_SparseHamiltonian.indices,
+                CSR_SparseHamiltonian.data,
+            )
+
+        if (
+            isinstance(measurementprocess.obs, (qml.ops.Hamiltonian, qml.Hermitian))
+            or (measurementprocess.obs.arithmetic_depth > 0)
+            or isinstance(measurementprocess.obs.name, List)
+        ):
+            ob_serialized = QuantumScriptSerializer(
+                self._qubit_state.device_name, self.dtype == np.complex64
+            )._ob(measurementprocess.obs)
+            return self._measurement_lightning.var(ob_serialized)
+
+        return self._measurement_lightning.var(
+            measurementprocess.obs.name, measurementprocess.obs.wires
+        )
+
+    def get_measurement_function(
+        self, measurementprocess: MeasurementProcess
+    ) -> Callable[[MeasurementProcess, TensorLike], TensorLike]:
+        """Get the appropriate method for performing a measurement.
+
+        Args:
+            measurementprocess (MeasurementProcess): measurement process to apply to the state
+
+        Returns:
+            Callable: function that returns the measurement result
+        """
+        if isinstance(measurementprocess, StateMeasurement):
+            if isinstance(measurementprocess, ExpectationMP):
+                if isinstance(measurementprocess.obs, (qml.Identity, qml.Projector)):
+                    return self.state_diagonalizing_gates
+                return self.expval
+
+            if isinstance(measurementprocess, ProbabilityMP):
+                return self.probs
+
+            if isinstance(measurementprocess, VarianceMP):
+                if isinstance(measurementprocess.obs, (qml.Identity, qml.Projector)):
+                    return self.state_diagonalizing_gates
+                return self.var
+            if measurementprocess.obs is None or measurementprocess.obs.has_diagonalizing_gates:
+                return self.state_diagonalizing_gates
+
+        raise NotImplementedError
+
+    def measurement(self, measurementprocess: MeasurementProcess) -> TensorLike:
+        """Apply a measurement process to a state.
+
+        Args:
+            measurementprocess (MeasurementProcess): measurement process to apply to the state
+
+        Returns:
+            TensorLike: the result of the measurement
+        """
+        return self.get_measurement_function(measurementprocess)(measurementprocess)
+
+    def measure_final_state(self, circuit: QuantumScript, mid_measurements=None) -> Result:
+        """
+        Perform the measurements required by the circuit on the provided state.
+
+        This is an internal function that will be called by the successor to ``lightning.kokkos``.
+
+        Args:
+            circuit (QuantumScript): The single circuit to simulate
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
+
+        Returns:
+            Tuple[TensorLike]: The measurement results
+        """
+
+        if not circuit.shots:
+            # analytic case
+            if len(circuit.measurements) == 1:
+                return self.measurement(circuit.measurements[0])
+
+            return tuple(self.measurement(mp) for mp in circuit.measurements)
+
+        # finite-shot case
+        results = self.measure_with_samples(
+            circuit.measurements,
+            shots=circuit.shots,
+            mid_measurements=mid_measurements,
+        )
+
+        if len(circuit.measurements) == 1:
+            if circuit.shots.has_partitioned_shots:
+                return tuple(res[0] for res in results)
+
+            return results[0]
+
+        return results
+
+    # pylint:disable = too-many-arguments
+    def measure_with_samples(
+        self,
+        measurements: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]],
+        shots: Shots,
+        mid_measurements=None,
+    ) -> List[TensorLike]:
+        """
+        Returns the samples of the measurement process performed on the given state.
+        This function assumes that the user-defined wire labels in the measurement process
+        have already been mapped to integer wires used in the device.
+
+        Args:
+            measurements (List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]]):
+                The sample measurements to perform
+            shots (Shots): The number of samples to take
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
+
+        Returns:
+            List[TensorLike[Any]]: Sample measurement results
+        """
+        # last N measurements are sampling MCMs in ``dynamic_one_shot`` execution mode
+        mps = measurements[0 : -len(mid_measurements)] if mid_measurements else measurements
+        groups, indices = _group_measurements(mps)
+
+        all_res = []
+        for group in groups:
+            if isinstance(group[0], (ExpectationMP, VarianceMP)) and isinstance(
+                group[0].obs, SparseHamiltonian
+            ):
+                raise TypeError(
+                    "ExpectationMP/VarianceMP(SparseHamiltonian) cannot be computed with samples."
+                )
+            if isinstance(group[0], VarianceMP) and isinstance(group[0].obs, (Hamiltonian, Sum)):
+                raise TypeError("VarianceMP(Hamiltonian/Sum) cannot be computed with samples.")
+            if isinstance(group[0], (ClassicalShadowMP, ShadowExpvalMP)):
+                raise TypeError(
+                    "ExpectationMP(ClassicalShadowMP, ShadowExpvalMP) cannot be computed with samples."
+                )
+            if isinstance(group[0], ExpectationMP) and isinstance(group[0].obs, Hamiltonian):
+                all_res.extend(self._measure_hamiltonian_with_samples(group, shots))
+            elif isinstance(group[0], ExpectationMP) and isinstance(group[0].obs, Sum):
+                all_res.extend(self._measure_sum_with_samples(group, shots))
+            else:
+                all_res.extend(self._measure_with_samples_diagonalizing_gates(group, shots))
+
+        # reorder results
+        flat_indices = []
+        for row in indices:
+            flat_indices += row
+        sorted_res = tuple(
+            res for _, res in sorted(list(enumerate(all_res)), key=lambda r: flat_indices[r[0]])
+        )
+
+        # append MCM samples
+        if mid_measurements:
+            sorted_res += tuple(mid_measurements.values())
+
+        # put the shot vector axis before the measurement axis
+        if shots.has_partitioned_shots:
+            sorted_res = tuple(zip(*sorted_res))
+
+        return sorted_res
+
+    def _apply_diagonalizing_gates(self, mps: List[SampleMeasurement], adjoint: bool = False):
+        if len(mps) == 1:
+            diagonalizing_gates = mps[0].diagonalizing_gates()
+        elif all(mp.obs for mp in mps):
+            diagonalizing_gates = qml.pauli.diagonalize_qwc_pauli_words([mp.obs for mp in mps])[0]
+        else:
+            diagonalizing_gates = []
+
+        if adjoint:
+            diagonalizing_gates = [
+                qml.adjoint(g, lazy=False) for g in reversed(diagonalizing_gates)
+            ]
+
+        self._qubit_state.apply_operations(diagonalizing_gates)
+
+    def _measure_with_samples_diagonalizing_gates(
+        self,
+        mps: List[SampleMeasurement],
+        shots: Shots,
+    ) -> TensorLike:
+        """
+        Returns the samples of the measurement process performed on the given state,
+        by rotating the state into the measurement basis using the diagonalizing gates
+        given by the measurement process.
+
+        Args:
+            mps (~.measurements.SampleMeasurement): The sample measurements to perform
+            shots (~.measurements.Shots): The number of samples to take
+
+        Returns:
+            TensorLike[Any]: Sample measurement results
+        """
+        # apply diagonalizing gates
+        self._apply_diagonalizing_gates(mps)
+
+        if self._mcmc:
+            total_indices = self._qubit_state.num_wires
+            wires = qml.wires.Wires(range(total_indices))
+        else:
+            #  ----------------------------------------
+            # Original:
+            # wires = reduce(sum, (mp.wires for mp in mps))
+            
+            # Specific for Kokkos:
+            total_indices = self._qubit_state.num_wires
+            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)
+
+        # 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(
+                    len(wires), self._kernel_name, self._num_burnin, shots.total_shots
+                ).astype(int, copy=False)
+            else:
+                #  ----------------------------------------
+                # Original:
+                # samples = self._measurement_lightning.generate_samples(
+                #     list(wires), shots.total_shots
+                # ).astype(int, copy=False)
+                
+                # Specific for Kokkos:
+                samples = self._measurement_lightning.generate_samples(
+                    len(wires), shots.total_shots
+                ).astype(int, copy=False)
+                #  ----------------------------------------
+                
+        except ValueError as e:
+            if str(e) != "probabilities contain NaN":
+                raise e
+            samples = qml.math.full((shots.total_shots, len(wires)), 0)
+
+        self._apply_diagonalizing_gates(mps, adjoint=True)
+
+        return _process_single_shot(samples)
+
+    def _measure_hamiltonian_with_samples(
+        self,
+        mp: List[SampleMeasurement],
+        shots: Shots,
+    ):
+        # the list contains only one element based on how we group measurements
+        mp = mp[0]
+
+        # if the measurement process involves a Hamiltonian, measure each
+        # of the terms separately and sum
+        def _sum_for_single_shot(s):
+            results = self.measure_with_samples(
+                [ExpectationMP(t) for t in mp.obs.terms()[1]],
+                s,
+            )
+            return sum(c * res for c, res in zip(mp.obs.terms()[0], results))
+
+        unsqueezed_results = tuple(_sum_for_single_shot(type(shots)(s)) for s in shots)
+        return [unsqueezed_results] if shots.has_partitioned_shots else [unsqueezed_results[0]]
+
+    def _measure_sum_with_samples(
+        self,
+        mp: List[SampleMeasurement],
+        shots: Shots,
+    ):
+        # the list contains only one element based on how we group measurements
+        mp = mp[0]
+
+        # if the measurement process involves a Sum, measure each
+        # of the terms separately and sum
+        def _sum_for_single_shot(s):
+            results = self.measure_with_samples(
+                [ExpectationMP(t) for t in mp.obs],
+                s,
+            )
+            return sum(results)
+
+        unsqueezed_results = tuple(_sum_for_single_shot(type(shots)(s)) for s in shots)
+        return [unsqueezed_results] if shots.has_partitioned_shots else [unsqueezed_results[0]]
diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index e54ff76fb..2d8bd8b0c 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -15,6 +15,33 @@
 Class implementation for lightning_kokkos state-vector manipulation.
 """
 
+from itertools import product
+
+import numpy as np
+import pennylane as qml
+from pennylane import BasisState, DeviceError, StatePrep
+from pennylane.measurements import MidMeasureMP
+from pennylane.ops import Conditional
+from pennylane.ops.op_math import Adjoint
+from pennylane.tape import QuantumScript
+from pennylane.wires import Wires
+
+try:
+    from pennylane_lightning.lightning_kokkos_ops import (
+        InitializationSettings,
+        StateVectorC64,
+        StateVectorC128,
+        allocate_aligned_array,
+        print_configuration,
+    )
+except ImportError:
+    pass  # Should be a complaint when kokkos_ops module is not available.
+
+from pennylane_lightning.core._serialize import global_phase_diagonal
+
+from ._measurements import LightningKokkosMeasurements
+
+
 
 class LightningKokkosStateVector:  # pylint: disable=too-few-public-methods
     """Lightning Kokkos state-vector class.
@@ -26,6 +53,454 @@ class LightningKokkosStateVector:  # pylint: disable=too-few-public-methods
         dtype: Datatypes for state-vector representation. Must be one of
             ``np.complex64`` or ``np.complex128``. Default is ``np.complex128``
         device_name(string): state vector device name. Options: ["lightning.kokkos"]
+        kokkos_args(InitializationSettings): binding for Kokkos::InitializationSettings
+            (threading parameters).
+        sync(bool): immediately sync with host-sv after applying operations
+
     """
 
-    pass  # pylint: disable=unnecessary-pass
+    def __init__(
+        self,
+        num_wires,
+        dtype=np.complex128,
+        device_name="lightning.kokkos",
+        kokkos_args=None,
+        sync=True,
+    ):  # pylint: disable=too-many-arguments
+        self._num_wires = num_wires
+        self._wires = Wires(range(num_wires))
+        self._dtype = dtype
+
+        self._kokkos_config = {}
+
+        if dtype not in [np.complex64, np.complex128]:  # pragma: no cover
+            raise TypeError(f"Unsupported complex type: {dtype}")
+
+        if device_name != "lightning.kokkos":
+            raise DeviceError(f'The device name "{device_name}" is not a valid option.')
+
+        self._device_name = device_name
+
+        # self._qubit_state = self._state_dtype()(self._num_wires)
+        if kokkos_args is None:
+            self._kokkos_state = self._state_dtype()(self.num_wires)
+        elif isinstance(kokkos_args, InitializationSettings):
+            self._kokkos_state = self._state_dtype()(self.num_wires, kokkos_args)
+        else:
+            type0 = type(InitializationSettings())
+            raise TypeError(
+                f"Argument kokkos_args must be of type {type0} but it is of {type(kokkos_args)}."
+            )
+        self._sync = sync
+
+        if not self._kokkos_config:
+            self._kokkos_config = self._kokkos_configuration()
+
+    @property
+    def dtype(self):
+        """Returns the state vector data type."""
+        return self._dtype
+
+    @property
+    def device_name(self):
+        """Returns the state vector device name."""
+        return self._device_name
+
+    @property
+    def wires(self):
+        """All wires that can be addressed on this device"""
+        return self._wires
+
+    @property
+    def num_wires(self):
+        """Number of wires addressed on this device"""
+        return self._num_wires
+
+    @property
+    def state_vector(self):
+        """Returns a handle to the state vector."""
+        return self._kokkos_state
+
+    @property
+    def state(self):
+        """Copy the state vector data from the device to the host.
+
+        A state vector Numpy array is explicitly allocated on the host to store and return
+        the data.
+
+        **Example**
+
+        >>> dev = qml.device('lightning.kokkos', wires=1)
+        >>> dev.apply([qml.PauliX(wires=[0])])
+        >>> print(dev.state)
+        [0.+0.j 1.+0.j]
+        """
+        state = np.zeros(2**self._num_wires, dtype=self.dtype)
+        state = self._asarray(state, dtype=self.dtype)
+        self.sync_d2h(state)
+        return state
+
+    @staticmethod
+    def _asarray(arr, dtype=None):
+        arr = np.asarray(arr)  # arr is not copied
+
+        if arr.dtype.kind not in ["f", "c"]:
+            return arr
+
+        if not dtype:
+            dtype = arr.dtype
+
+        # We allocate a new aligned memory and copy data to there if alignment
+        # or dtype mismatches
+        # Note that get_alignment does not necessarily return CPUMemoryModel(Unaligned) even for
+        # numpy allocated memory as the memory location happens to be aligned.
+        if arr.dtype != dtype:
+            new_arr = allocate_aligned_array(arr.size, np.dtype(dtype), False).reshape(arr.shape)
+            np.copyto(new_arr, arr)
+            arr = new_arr
+        return arr
+
+    def sync_h2d(self, state_vector):
+        """Copy the state vector data on host provided by the user to the state
+        vector on the device
+
+        Args:
+            state_vector(array[complex]): the state vector array on host.
+
+
+        **Example**
+
+        >>> dev = qml.device('lightning.kokkos', wires=3)
+        >>> obs = qml.Identity(0) @ qml.PauliX(1) @ qml.PauliY(2)
+        >>> obs1 = qml.Identity(1)
+        >>> H = qml.Hamiltonian([1.0, 1.0], [obs1, obs])
+        >>> state_vector = np.array([0.0 + 0.0j, 0.0 + 0.1j, 0.1 + 0.1j, 0.1 + 0.2j, 0.2 + 0.2j, 0.3 + 0.3j, 0.3 + 0.4j, 0.4 + 0.5j,], dtype=np.complex64)
+        >>> dev.sync_h2d(state_vector)
+        >>> res = dev.expval(H)
+        >>> print(res)
+        1.0
+        """
+        self._kokkos_state.HostToDevice(state_vector.ravel(order="C"))
+
+    def sync_d2h(self, state_vector):
+        """Copy the state vector data on device to a state vector on the host provided
+        by the user
+
+        Args:
+            state_vector(array[complex]): the state vector array on host
+
+
+        **Example**
+
+        >>> dev = qml.device('lightning.kokkos', wires=1)
+        >>> dev.apply([qml.PauliX(wires=[0])])
+        >>> state_vector = np.zeros(2**dev.num_wires).astype(dev.C_DTYPE)
+        >>> dev.sync_d2h(state_vector)
+        >>> print(state_vector)
+        [0.+0.j 1.+0.j]
+        """
+        self._kokkos_state.DeviceToHost(state_vector.ravel(order="C"))
+
+    def _kokkos_configuration(self):
+        """Set the default configuration of the kokkos device.
+        Returns: kokkos configuration
+        """
+        return print_configuration()
+
+    def _state_dtype(self):
+        """Binding to Lightning Managed state vector C++ class.
+
+        Returns: the state vector class
+        """
+        return StateVectorC128 if self.dtype == np.complex128 else StateVectorC64
+
+    def _create_basis_state(self, index):
+        """Return a computational basis state over all wires.
+
+        Args:
+            index (int): integer representing the computational basis state.
+        """
+        self._kokkos_state.setBasisState(index)
+
+    def reset_state(self):
+        """Reset the device's state"""
+        # init the state vector to |00..0>
+        self._kokkos_state.resetStateVector()
+
+    def _preprocess_state_vector(self, state, device_wires):
+        """Initialize the internal state vector in a specified state.
+
+        Args:
+            state (array[complex]): normalized input state of length ``2**len(wires)``
+                or broadcasted state of shape ``(batch_size, 2**len(wires))``
+            device_wires (Wires): wires that get initialized in the state
+
+        Returns:
+            array[int]: indices for which the state is changed to input state vector elements
+            array[complex]: normalized input state of length ``2**len(wires)``
+                or broadcasted state of shape ``(batch_size, 2**len(wires))``
+        """
+        # special case for integral types
+        if state.dtype.kind == "i":
+            state = np.array(state, dtype=self.dtype)
+
+        if len(device_wires) == self._num_wires and Wires(sorted(device_wires)) == device_wires:
+            return None, state
+
+        # generate basis states on subset of qubits via the cartesian product
+        basis_states = np.array(list(product([0, 1], repeat=len(device_wires))))
+
+        # get basis states to alter on full set of qubits
+        unravelled_indices = np.zeros((2 ** len(device_wires), self._num_wires), dtype=int)
+        unravelled_indices[:, device_wires] = basis_states
+
+        # get indices for which the state is changed to input state vector elements
+        ravelled_indices = np.ravel_multi_index(unravelled_indices.T, [2] * self._num_wires)
+        return ravelled_indices, state
+
+    def _get_basis_state_index(self, state, wires):
+        """Returns the basis state index of a specified computational basis state.
+
+        Args:
+            state (array[int]): computational basis state of shape ``(wires,)``
+                consisting of 0s and 1s
+            wires (Wires): wires that the provided computational state should be initialized on
+
+        Returns:
+            int: basis state index
+        """
+        # length of basis state parameter
+        n_basis_state = len(state)
+
+        if not set(state.tolist()).issubset({0, 1}):
+            raise ValueError("BasisState parameter must consist of 0 or 1 integers.")
+
+        if n_basis_state != len(wires):
+            raise ValueError("BasisState parameter and wires must be of equal length.")
+
+        # get computational basis state number
+        basis_states = 2 ** (self._num_wires - 1 - np.array(wires))
+        basis_states = qml.math.convert_like(basis_states, state)
+        return int(qml.math.dot(state, basis_states))
+
+    def _apply_state_vector(self, state, device_wires: Wires):
+        """Initialize the internal state vector in a specified state.
+        Args:
+            state (array[complex]): normalized input state of length ``2**len(wires)``
+                or broadcasted state of shape ``(batch_size, 2**len(wires))``
+            device_wires (Wires): wires that get initialized in the state
+        """
+
+        if isinstance(state, self._kokkos_state.__class__):
+            state_data = allocate_aligned_array(state.size, np.dtype(self.dtype), True)
+            state.DeviceToHost(state_data)
+            state = state_data
+
+        ravelled_indices, state = self._preprocess_state_vector(state, device_wires)
+
+        # translate to wire labels used by device
+        output_shape = [2] * self._num_wires
+
+        if len(device_wires) == self._num_wires and Wires(sorted(device_wires)) == device_wires:
+            # Initialize the entire device state with the input state
+            self.sync_h2d(np.reshape(state, output_shape))
+            return
+
+        self._kokkos_state.setStateVector(ravelled_indices, state)  # this operation on device
+
+    def _apply_basis_state(self, state, wires):
+        """Initialize the state vector in a specified computational basis state.
+
+        Args:
+            state (array[int]): computational basis state of shape ``(wires,)``
+                consisting of 0s and 1s.
+            wires (Wires): wires that the provided computational state should be
+                initialized on
+
+        Note: This function does not support broadcasted inputs yet.
+        """
+        num = self._get_basis_state_index(state, wires)
+        self._create_basis_state(num)
+
+    def _apply_lightning_controlled(self, operation):
+        """Apply an arbitrary controlled operation to the state tensor.
+
+        Args:
+            operation (~pennylane.operation.Operation): controlled operation to apply
+
+        Returns:
+            None
+        """
+        state = self.state_vector
+
+        basename = operation.base.name
+        method = getattr(state, f"{basename}", None)
+        control_wires = list(operation.control_wires)
+        control_values = operation.control_values
+        #  ----------------------------------------
+        # Original:
+        # target_wires = list(operation.target_wires)
+
+        # Specific for Kokkos:
+        name = operation.name
+        #  ----------------------------------------
+        if method is not None:  # apply n-controlled specialized gate
+            inv = False
+            #  ----------------------------------------
+            # Original:
+            # param = operation.parameters
+            # method(control_wires, control_values, target_wires, inv, param)
+
+            # Specific for Kokkos:
+            param = operation.parameters[0]
+            wires = self.wires.indices(operation.wires)
+            matrix = global_phase_diagonal(param, self.wires, control_wires, control_values)
+            state.apply(name, wires, inv, [[param]], matrix)
+            #  ----------------------------------------
+
+        else:  # apply gate as an n-controlled matrix
+            #  ----------------------------------------
+            # Original:
+            # method = getattr(state, "applyControlledMatrix")
+            # method(
+            #     qml.matrix(operation.base),
+            #     control_wires,
+            #     control_values,
+            #     target_wires,
+            #     False,
+            # )
+
+            # Specific for Kokkos:
+            raise ValueError("Unsupported apply Controlled Matrix")
+            #  ----------------------------------------
+
+    def _apply_lightning_midmeasure(
+        self, operation: MidMeasureMP, mid_measurements: dict, postselect_mode: str
+    ):
+        """Execute a MidMeasureMP operation and return the sample in mid_measurements.
+
+        Args:
+            operation (~pennylane.operation.Operation): mid-circuit measurement
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
+            postselect_mode (str): Configuration for handling shots with mid-circuit measurement
+                postselection. Use ``"hw-like"`` to discard invalid shots and ``"fill-shots"`` to
+                keep the same number of shots.
+
+        Returns:
+            None
+        """
+        wires = self.wires.indices(operation.wires)
+        wire = list(wires)[0]
+        circuit = QuantumScript([], [qml.sample(wires=operation.wires)], shots=1)
+        if postselect_mode == "fill-shots" and operation.postselect is not None:
+            sample = operation.postselect
+        else:
+            sample = LightningKokkosMeasurements(self).measure_final_state(circuit)
+            sample = np.squeeze(sample)
+        mid_measurements[operation] = sample
+        getattr(self.state_vector, "collapse")(wire, bool(sample))
+        if operation.reset and bool(sample):
+            self.apply_operations([qml.PauliX(operation.wires)], mid_measurements=mid_measurements)
+
+    def _apply_lightning(
+        self, operations, mid_measurements: dict = None, postselect_mode: str = None
+    ):
+        """Apply a list of operations to the state tensor.
+
+        Args:
+            operations (list[~pennylane.operation.Operation]): operations to apply
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
+            postselect_mode (str): Configuration for handling shots with mid-circuit measurement
+                postselection. Use ``"hw-like"`` to discard invalid shots and ``"fill-shots"`` to
+                keep the same number of shots. Default is ``None``.
+
+        Returns:
+            None
+        """
+        state = self.state_vector
+
+        # Skip over identity operations instead of performing
+        # matrix multiplication with it.
+        for operation in operations:
+            if isinstance(operation, qml.Identity):
+                continue
+            if isinstance(operation, Adjoint):
+                name = operation.base.name
+                invert_param = True
+            else:
+                name = operation.name
+                invert_param = False
+            method = getattr(state, name, None)
+            wires = list(operation.wires)
+
+            if isinstance(operation, Conditional):
+                if operation.meas_val.concretize(mid_measurements):
+                    self._apply_lightning([operation.base])
+            elif isinstance(operation, MidMeasureMP):
+                self._apply_lightning_midmeasure(
+                    operation, mid_measurements, postselect_mode=postselect_mode
+                )
+            elif method is not None:  # apply specialized gate
+                param = operation.parameters
+                method(wires, invert_param, param)
+            elif (
+                isinstance(operation, qml.ops.Controlled)
+                and isinstance(operation.base, qml.GlobalPhase)
+                # Specific for Kokkos:
+                # Kokkos do not support the controlled gates except for GlobalPhase
+            ):  # apply n-controlled gate
+                self._apply_lightning_controlled(operation)
+            else:  # apply gate as a matrix
+                # Inverse can be set to False since qml.matrix(operation) is already in
+                # inverted form
+                method = getattr(state, "applyMatrix")
+                try:
+                    method(qml.matrix(operation), wires, False)
+                except AttributeError:  # pragma: no cover
+                    # To support older versions of PL
+                    method(operation.matrix, wires, False)
+
+    def apply_operations(
+        self, operations, mid_measurements: dict = None, postselect_mode: str = None
+    ):
+        """Applies operations to the state vector."""
+        # State preparation is currently done in Python
+        if operations:  # make sure operations[0] exists
+            if isinstance(operations[0], StatePrep):
+                self._apply_state_vector(operations[0].parameters[0].copy(), operations[0].wires)
+                operations = operations[1:]
+            elif isinstance(operations[0], BasisState):
+                self._apply_basis_state(operations[0].parameters[0], operations[0].wires)
+                operations = operations[1:]
+
+        self._apply_lightning(
+            operations, mid_measurements=mid_measurements, postselect_mode=postselect_mode
+        )
+
+    def get_final_state(
+        self,
+        circuit: QuantumScript,
+        mid_measurements: dict = None,
+        postselect_mode: str = None,
+    ):
+        """
+        Get the final state that results from executing the given quantum script.
+
+        This is an internal function that will be called by the successor to ``lightning.qubit``.
+
+        Args:
+            circuit (QuantumScript): The single circuit to simulate
+            mid_measurements (None, dict): Dictionary of mid-circuit measurements
+            postselect_mode (str): Configuration for handling shots with mid-circuit measurement
+                postselection. Use ``"hw-like"`` to discard invalid shots and ``"fill-shots"`` to
+                keep the same number of shots. Default is ``None``.
+
+        Returns:
+            LightningStateVector: Lightning final state class.
+
+        """
+        self.apply_operations(
+            circuit.operations, mid_measurements=mid_measurements, postselect_mode=postselect_mode
+        )
+
+        return self
diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
index e91722afb..f6190dbc9 100644
--- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
+++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
@@ -26,8 +26,10 @@
 from pennylane.devices.modifiers import simulator_tracking, single_tape_support
 from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.typing import Result, ResultBatch
+from pennylane.measurements import MidMeasureMP
 
 from ._state_vector import LightningKokkosStateVector
+from ._measurements import LightningKokkosMeasurements
 
 try:
     # pylint: disable=import-error, no-name-in-module
@@ -44,9 +46,10 @@
 PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
 
 
-def simulate(  # pylint: disable=unused-argument
+def simulate(
     circuit: QuantumScript,
     state: LightningKokkosStateVector,
+    # mcmc: dict = None, # L-Kokkos has no support for MCMC sampling
     postselect_mode: str = None,
 ) -> Result:
     """Simulate a single quantum script.
@@ -63,7 +66,33 @@ def simulate(  # pylint: disable=unused-argument
 
     Note that this function can return measurements for non-commuting observables simultaneously.
     """
-    return 0
+    if mcmc is None:
+        mcmc = {}
+    state.reset_state()
+    has_mcm = any(isinstance(op, MidMeasureMP) for op in circuit.operations)
+    if circuit.shots and has_mcm:
+        results = []
+        aux_circ = qml.tape.QuantumScript(
+            circuit.operations,
+            circuit.measurements,
+            shots=[1],
+            trainable_params=circuit.trainable_params,
+        )
+        for _ in range(circuit.shots.total_shots):
+            state.reset_state()
+            mid_measurements = {}
+            final_state = state.get_final_state(
+                aux_circ, mid_measurements=mid_measurements, postselect_mode=postselect_mode
+            )
+            results.append(
+                LightningKokkosMeasurements(final_state, **mcmc).measure_final_state(
+                    aux_circ, mid_measurements=mid_measurements
+                )
+            )
+        return tuple(results)
+    
+    final_state = state.get_final_state(circuit)
+    return LightningKokkosMeasurements(final_state, **mcmc).measure_final_state(circuit)
 
 
 def jacobian(  # pylint: disable=unused-argument
@@ -320,9 +349,14 @@ def __init__(  # pylint: disable=too-many-arguments
 
         self._statevector = LightningKokkosStateVector(num_wires=len(self.wires), dtype=c_dtype)
 
+        #  ----------------------------------------
+        # Original:
         # TODO: Investigate usefulness of creating numpy random generator
-        seed = np.random.randint(0, high=10000000) if seed == "global" else seed
-        self._rng = np.random.default_rng(seed)
+        # seed = np.random.randint(0, high=10000000) if seed == "global" else seed
+        # self._rng = np.random.default_rng(seed)
+        
+        # Specific for Kokkos:
+        #  ----------------------------------------        
 
         self._c_dtype = c_dtype
         self._batch_obs = batch_obs
diff --git a/tests/lightning_qubit/test_measurements_class.py b/tests/lightning_qubit/test_measurements_class.py
index 7fc142973..984ef5531 100644
--- a/tests/lightning_qubit/test_measurements_class.py
+++ b/tests/lightning_qubit/test_measurements_class.py
@@ -19,7 +19,7 @@
 import numpy as np
 import pennylane as qml
 import pytest
-from conftest import LightningDevice, device_name  # tested device
+from conftest import THETA, PHI, LightningDevice, device_name  # tested device
 from flaky import flaky
 from pennylane.devices import DefaultQubit
 from pennylane.measurements import VarianceMP
@@ -30,21 +30,21 @@
 except ImportError:
     pass
 
-from pennylane_lightning.lightning_qubit._measurements import LightningMeasurements
-from pennylane_lightning.lightning_qubit._state_vector import LightningStateVector
+if device_name == 'lightning.qubit':
+    from pennylane_lightning.lightning_qubit._measurements import LightningMeasurements
+    from pennylane_lightning.lightning_qubit._state_vector import LightningStateVector
+    
+if device_name == 'lightning.kokkos':
+    from pennylane_lightning.lightning_kokkos._measurements import LightningKokkosMeasurements as LightningMeasurements
+    from pennylane_lightning.lightning_kokkos._state_vector import LightningKokkosStateVector as LightningStateVector
 
-if device_name == "lightning.kokkos":
-    pytest.skip("Kokkos new API in WIP.  Skipping.", allow_module_level=True)
+if device_name not in ("lightning.qubit", "lightning.kokkos"):
+    pytest.skip("Exclusive tests for lightning.qubit or lightning.kokkos. Skipping.", allow_module_level=True)
 
-if device_name != "lightning.qubit":
-    pytest.skip("Exclusive tests for lightning.qubit. Skipping.", allow_module_level=True)
 
 if not LightningDevice._CPP_BINARY_AVAILABLE:
     pytest.skip("No binary module found. Skipping.", allow_module_level=True)
 
-THETA = np.linspace(0.11, 1, 3)
-PHI = np.linspace(0.32, 1, 3)
-
 
 # General LightningStateVector fixture, for any number of wires.
 @pytest.fixture(
@@ -584,6 +584,10 @@ def test_double_return_value(self, shots, measurement, obs0_, obs1_, lightning_s
             [[1, 0], [0, 1]],
         ],
     )
+    @pytest.mark.skipif(
+        device_name == "lightning.kokkos",
+        reason="Kokkos new API in WIP.  Skipping.",
+    )
     def test_probs_tape_unordered_wires(self, cases, tol):
         """Test probs with a circuit on wires=[0] fails for out-of-order wires passed to probs."""
 
diff --git a/tests/lightning_qubit/test_state_vector_class.py b/tests/lightning_qubit/test_state_vector_class.py
index 0cc8fe303..0b461017b 100644
--- a/tests/lightning_qubit/test_state_vector_class.py
+++ b/tests/lightning_qubit/test_state_vector_class.py
@@ -24,13 +24,15 @@
 from pennylane.tape import QuantumScript
 from pennylane.wires import Wires
 
-from pennylane_lightning.lightning_qubit._state_vector import LightningStateVector
+if device_name == 'lightning.qubit':
+    from pennylane_lightning.lightning_qubit._state_vector import LightningStateVector
+    
+if device_name == 'lightning.kokkos':
+    from pennylane_lightning.lightning_kokkos._state_vector import LightningKokkosStateVector as LightningStateVector
 
-if device_name == "lightning.kokkos":
-    pytest.skip("Kokkos new API in WIP.  Skipping.", allow_module_level=True)
 
-if device_name != "lightning.qubit":
-    pytest.skip("Exclusive tests for lightning.qubit. Skipping.", allow_module_level=True)
+if device_name not in ("lightning.qubit", "lightning.kokkos"):
+    pytest.skip("Exclusive tests for lightning.qubit or lightning.kokkos. Skipping.", allow_module_level=True)
 
 if not LightningDevice._CPP_BINARY_AVAILABLE:
     pytest.skip("No binary module found. Skipping.", allow_module_level=True)
@@ -38,7 +40,8 @@
 
 @pytest.mark.parametrize("num_wires", range(4))
 @pytest.mark.parametrize("dtype", [np.complex64, np.complex128])
-@pytest.mark.parametrize("device_name", ["lightning.qubit"])
+# @pytest.mark.parametrize("device_name", ["lightning.qubit"])
+@pytest.mark.parametrize("device_name", [device_name])
 def test_device_name_and_init(num_wires, dtype, device_name):
     """Test the class initialization and returned properties."""
     state_vector = LightningStateVector(num_wires, dtype=dtype, device_name=device_name)

From 145e18e7090a0e07f24c933f7fcd188457c0f482 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Thu, 25 Jul 2024 15:38:09 -0400
Subject: [PATCH 02/26] Apply format

---
 .../lightning_kokkos/_measurements.py         |  6 +++---
 .../lightning_kokkos/_state_vector.py         |  1 -
 .../lightning_kokkos/lightning_kokkos.py      | 16 ++++++++------
 .../test_measurements_class.py                | 21 ++++++++++++-------
 .../test_state_vector_class.py                | 15 ++++++++-----
 5 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py
index fc331a799..fcb9d8924 100644
--- a/pennylane_lightning/lightning_kokkos/_measurements.py
+++ b/pennylane_lightning/lightning_kokkos/_measurements.py
@@ -391,7 +391,7 @@ def _measure_with_samples_diagonalizing_gates(
             #  ----------------------------------------
             # Original:
             # wires = reduce(sum, (mp.wires for mp in mps))
-            
+
             # Specific for Kokkos:
             total_indices = self._qubit_state.num_wires
             wires = qml.wires.Wires(range(total_indices))
@@ -444,13 +444,13 @@ def _process_single_shot(samples):
                 # samples = self._measurement_lightning.generate_samples(
                 #     list(wires), shots.total_shots
                 # ).astype(int, copy=False)
-                
+
                 # Specific for Kokkos:
                 samples = self._measurement_lightning.generate_samples(
                     len(wires), shots.total_shots
                 ).astype(int, copy=False)
                 #  ----------------------------------------
-                
+
         except ValueError as e:
             if str(e) != "probabilities contain NaN":
                 raise e
diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index 2d8bd8b0c..b203b3fd7 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -42,7 +42,6 @@
 from ._measurements import LightningKokkosMeasurements
 
 
-
 class LightningKokkosStateVector:  # pylint: disable=too-few-public-methods
     """Lightning Kokkos state-vector class.
 
diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
index f6190dbc9..cb10e4315 100644
--- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
+++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
@@ -24,12 +24,12 @@
 import pennylane as qml
 from pennylane.devices import DefaultExecutionConfig, Device, ExecutionConfig
 from pennylane.devices.modifiers import simulator_tracking, single_tape_support
+from pennylane.measurements import MidMeasureMP
 from pennylane.tape import QuantumScript, QuantumTape
 from pennylane.typing import Result, ResultBatch
-from pennylane.measurements import MidMeasureMP
 
-from ._state_vector import LightningKokkosStateVector
 from ._measurements import LightningKokkosMeasurements
+from ._state_vector import LightningKokkosStateVector
 
 try:
     # pylint: disable=import-error, no-name-in-module
@@ -90,7 +90,7 @@ def simulate(
                 )
             )
         return tuple(results)
-    
+
     final_state = state.get_final_state(circuit)
     return LightningKokkosMeasurements(final_state, **mcmc).measure_final_state(circuit)
 
@@ -328,6 +328,10 @@ def __init__(  # pylint: disable=too-many-arguments
         *,
         c_dtype=np.complex128,
         shots=None,
+        # seed="global", Specific for LQubit
+        # mcmc=False, Specific for LQubit
+        # kernel_name="Local", Specific for LQubit
+        # num_burnin=100, Specific for LQubit
         batch_obs=False,
         # Kokkos arguments
         sync=True,
@@ -354,9 +358,9 @@ def __init__(  # pylint: disable=too-many-arguments
         # TODO: Investigate usefulness of creating numpy random generator
         # seed = np.random.randint(0, high=10000000) if seed == "global" else seed
         # self._rng = np.random.default_rng(seed)
-        
-        # Specific for Kokkos:
-        #  ----------------------------------------        
+
+        # Specific for LKokkos:
+        #  ----------------------------------------
 
         self._c_dtype = c_dtype
         self._batch_obs = batch_obs
diff --git a/tests/lightning_qubit/test_measurements_class.py b/tests/lightning_qubit/test_measurements_class.py
index 984ef5531..33aad93f9 100644
--- a/tests/lightning_qubit/test_measurements_class.py
+++ b/tests/lightning_qubit/test_measurements_class.py
@@ -19,7 +19,7 @@
 import numpy as np
 import pennylane as qml
 import pytest
-from conftest import THETA, PHI, LightningDevice, device_name  # tested device
+from conftest import PHI, THETA, LightningDevice, device_name  # tested device
 from flaky import flaky
 from pennylane.devices import DefaultQubit
 from pennylane.measurements import VarianceMP
@@ -30,16 +30,23 @@
 except ImportError:
     pass
 
-if device_name == 'lightning.qubit':
+if device_name == "lightning.qubit":
     from pennylane_lightning.lightning_qubit._measurements import LightningMeasurements
     from pennylane_lightning.lightning_qubit._state_vector import LightningStateVector
-    
-if device_name == 'lightning.kokkos':
-    from pennylane_lightning.lightning_kokkos._measurements import LightningKokkosMeasurements as LightningMeasurements
-    from pennylane_lightning.lightning_kokkos._state_vector import LightningKokkosStateVector as LightningStateVector
+
+if device_name == "lightning.kokkos":
+    from pennylane_lightning.lightning_kokkos._measurements import (
+        LightningKokkosMeasurements as LightningMeasurements,
+    )
+    from pennylane_lightning.lightning_kokkos._state_vector import (
+        LightningKokkosStateVector as LightningStateVector,
+    )
 
 if device_name not in ("lightning.qubit", "lightning.kokkos"):
-    pytest.skip("Exclusive tests for lightning.qubit or lightning.kokkos. Skipping.", allow_module_level=True)
+    pytest.skip(
+        "Exclusive tests for lightning.qubit or lightning.kokkos. Skipping.",
+        allow_module_level=True,
+    )
 
 
 if not LightningDevice._CPP_BINARY_AVAILABLE:
diff --git a/tests/lightning_qubit/test_state_vector_class.py b/tests/lightning_qubit/test_state_vector_class.py
index 0b461017b..5cea7c3e0 100644
--- a/tests/lightning_qubit/test_state_vector_class.py
+++ b/tests/lightning_qubit/test_state_vector_class.py
@@ -24,15 +24,20 @@
 from pennylane.tape import QuantumScript
 from pennylane.wires import Wires
 
-if device_name == 'lightning.qubit':
+if device_name == "lightning.qubit":
     from pennylane_lightning.lightning_qubit._state_vector import LightningStateVector
-    
-if device_name == 'lightning.kokkos':
-    from pennylane_lightning.lightning_kokkos._state_vector import LightningKokkosStateVector as LightningStateVector
+
+if device_name == "lightning.kokkos":
+    from pennylane_lightning.lightning_kokkos._state_vector import (
+        LightningKokkosStateVector as LightningStateVector,
+    )
 
 
 if device_name not in ("lightning.qubit", "lightning.kokkos"):
-    pytest.skip("Exclusive tests for lightning.qubit or lightning.kokkos. Skipping.", allow_module_level=True)
+    pytest.skip(
+        "Exclusive tests for lightning.qubit or lightning.kokkos. Skipping.",
+        allow_module_level=True,
+    )
 
 if not LightningDevice._CPP_BINARY_AVAILABLE:
     pytest.skip("No binary module found. Skipping.", allow_module_level=True)

From a5ae5a7f0c310f31081943bb37ba6b37d21c9db0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Thu, 25 Jul 2024 17:19:58 -0400
Subject: [PATCH 03/26] Remove mcmc from functions

---
 .../lightning_kokkos/_measurements.py         | 102 ++++++++----------
 .../lightning_kokkos/lightning_kokkos.py      |  20 +---
 2 files changed, 47 insertions(+), 75 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py
index fcb9d8924..3f274fbf5 100644
--- a/pennylane_lightning/lightning_kokkos/_measurements.py
+++ b/pennylane_lightning/lightning_kokkos/_measurements.py
@@ -54,35 +54,15 @@ class LightningKokkosMeasurements:
 
     Args:
         qubit_state(LightningStateVector): Lightning state-vector class containing the state vector to be measured.
-        mcmc (bool): Determine whether to use the approximate Markov Chain Monte Carlo
-            sampling method when generating samples.
-        kernel_name (str): name of MCMC transition kernel. The current version supports
-            two kernels: ``"Local"`` and ``"NonZeroRandom"``.
-            The local kernel conducts a bit-flip local transition between states.
-            The local kernel generates a random qubit site and then generates a random
-            number to determine the new bit at that qubit site. The ``"NonZeroRandom"`` kernel
-            randomly transits between states that have nonzero probability.
-        num_burnin (int): number of MCMC steps that will be dropped. Increasing this value will
-            result in a closer approximation but increased runtime.
     """
 
     def __init__(
         self,
         kokkos_state,
-        mcmc: bool = None,
-        kernel_name: str = None,
-        num_burnin: int = None,
     ) -> None:
         self._qubit_state = kokkos_state
         self._dtype = kokkos_state.dtype
         self._measurement_lightning = self._measurement_dtype()(kokkos_state.state_vector)
-        self._mcmc = mcmc
-        self._kernel_name = kernel_name
-        self._num_burnin = num_burnin
-        if self._mcmc and not self._kernel_name:
-            self._kernel_name = "Local"
-        if self._mcmc and not self._num_burnin:
-            self._num_burnin = 100
 
     @property
     def qubit_state(self):
@@ -384,18 +364,18 @@ def _measure_with_samples_diagonalizing_gates(
         # apply diagonalizing gates
         self._apply_diagonalizing_gates(mps)
 
-        if self._mcmc:
-            total_indices = self._qubit_state.num_wires
-            wires = qml.wires.Wires(range(total_indices))
-        else:
-            #  ----------------------------------------
-            # Original:
-            # wires = reduce(sum, (mp.wires for mp in mps))
-
-            # Specific for Kokkos:
-            total_indices = self._qubit_state.num_wires
-            wires = qml.wires.Wires(range(total_indices))
-            #  ----------------------------------------
+        #  ----------------------------------------
+        # Original:
+        # if self._mcmc:
+        #     total_indices = self._qubit_state.num_wires
+        #     wires = qml.wires.Wires(range(total_indices))
+        # else:
+        #     wires = reduce(sum, (mp.wires for mp in mps))
+            
+        # Specific for Kokkos:
+        total_indices = self._qubit_state.num_wires
+        wires = qml.wires.Wires(range(total_indices))
+        #  ----------------------------------------
 
         def _process_single_shot(samples):
             processed = []
@@ -416,14 +396,22 @@ def _process_single_shot(samples):
                 # 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)
+                    #  ----------------------------------------
+                    # 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
@@ -434,23 +422,23 @@ def _process_single_shot(samples):
             return tuple(zip(*processed_samples))
 
         try:
-            if self._mcmc:
-                samples = self._measurement_lightning.generate_mcmc_samples(
-                    len(wires), self._kernel_name, self._num_burnin, shots.total_shots
-                ).astype(int, copy=False)
-            else:
-                #  ----------------------------------------
-                # Original:
-                # samples = self._measurement_lightning.generate_samples(
-                #     list(wires), shots.total_shots
-                # ).astype(int, copy=False)
-
-                # Specific for Kokkos:
-                samples = self._measurement_lightning.generate_samples(
-                    len(wires), shots.total_shots
-                ).astype(int, copy=False)
-                #  ----------------------------------------
-
+            #  ----------------------------------------
+            # Original:
+            # if self._mcmc:
+            #     samples = self._measurement_lightning.generate_mcmc_samples(
+            #         len(wires), self._kernel_name, self._num_burnin, shots.total_shots
+            #     ).astype(int, copy=False)
+            # else:
+            #     samples = self._measurement_lightning.generate_samples(
+            #         list(wires), shots.total_shots
+            #     ).astype(int, copy=False)
+                
+            # Specific for Kokkos:
+            samples = self._measurement_lightning.generate_samples(
+                len(wires), shots.total_shots
+            ).astype(int, copy=False)
+            #  ----------------------------------------
+                
         except ValueError as e:
             if str(e) != "probabilities contain NaN":
                 raise e
diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
index cb10e4315..5b21fd3ed 100644
--- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
+++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
@@ -49,7 +49,6 @@
 def simulate(
     circuit: QuantumScript,
     state: LightningKokkosStateVector,
-    # mcmc: dict = None, # L-Kokkos has no support for MCMC sampling
     postselect_mode: str = None,
 ) -> Result:
     """Simulate a single quantum script.
@@ -66,8 +65,6 @@ def simulate(
 
     Note that this function can return measurements for non-commuting observables simultaneously.
     """
-    if mcmc is None:
-        mcmc = {}
     state.reset_state()
     has_mcm = any(isinstance(op, MidMeasureMP) for op in circuit.operations)
     if circuit.shots and has_mcm:
@@ -85,14 +82,14 @@ def simulate(
                 aux_circ, mid_measurements=mid_measurements, postselect_mode=postselect_mode
             )
             results.append(
-                LightningKokkosMeasurements(final_state, **mcmc).measure_final_state(
+                LightningKokkosMeasurements(final_state).measure_final_state(
                     aux_circ, mid_measurements=mid_measurements
                 )
             )
         return tuple(results)
 
     final_state = state.get_final_state(circuit)
-    return LightningKokkosMeasurements(final_state, **mcmc).measure_final_state(circuit)
+    return LightningKokkosMeasurements(final_state).measure_final_state(circuit)
 
 
 def jacobian(  # pylint: disable=unused-argument
@@ -328,10 +325,6 @@ def __init__(  # pylint: disable=too-many-arguments
         *,
         c_dtype=np.complex128,
         shots=None,
-        # seed="global", Specific for LQubit
-        # mcmc=False, Specific for LQubit
-        # kernel_name="Local", Specific for LQubit
-        # num_burnin=100, Specific for LQubit
         batch_obs=False,
         # Kokkos arguments
         sync=True,
@@ -353,15 +346,6 @@ def __init__(  # pylint: disable=too-many-arguments
 
         self._statevector = LightningKokkosStateVector(num_wires=len(self.wires), dtype=c_dtype)
 
-        #  ----------------------------------------
-        # Original:
-        # TODO: Investigate usefulness of creating numpy random generator
-        # seed = np.random.randint(0, high=10000000) if seed == "global" else seed
-        # self._rng = np.random.default_rng(seed)
-
-        # Specific for LKokkos:
-        #  ----------------------------------------
-
         self._c_dtype = c_dtype
         self._batch_obs = batch_obs
 

From 97d0182774895bc8145a4223c500c90bc48aff13 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 08:25:34 -0400
Subject: [PATCH 04/26] apply format

---
 .../lightning_kokkos/_measurements.py              | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py
index 3f274fbf5..f641c0fdd 100644
--- a/pennylane_lightning/lightning_kokkos/_measurements.py
+++ b/pennylane_lightning/lightning_kokkos/_measurements.py
@@ -371,7 +371,7 @@ def _measure_with_samples_diagonalizing_gates(
         #     wires = qml.wires.Wires(range(total_indices))
         # else:
         #     wires = reduce(sum, (mp.wires for mp in mps))
-            
+
         # Specific for Kokkos:
         total_indices = self._qubit_state.num_wires
         wires = qml.wires.Wires(range(total_indices))
@@ -406,11 +406,11 @@ def _process_single_shot(samples):
                     #     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)
+                    samples = self._measurement_lightning.generate_samples(len(wires), s).astype(
+                        int, copy=False
+                    )
                     #  ----------------------------------------
                 except ValueError as e:
                     if str(e) != "probabilities contain NaN":
@@ -432,13 +432,13 @@ def _process_single_shot(samples):
             #     samples = self._measurement_lightning.generate_samples(
             #         list(wires), shots.total_shots
             #     ).astype(int, copy=False)
-                
+
             # Specific for Kokkos:
             samples = self._measurement_lightning.generate_samples(
                 len(wires), shots.total_shots
             ).astype(int, copy=False)
             #  ----------------------------------------
-                
+
         except ValueError as e:
             if str(e) != "probabilities contain NaN":
                 raise e

From fe502eaab72765cb52e96a6d9cea5fe2f8ec4903 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 08:25:44 -0400
Subject: [PATCH 05/26] trigger CIs


From 58a064cf000bf4303de40f4009107c9d8ba3ebb2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 10:17:32 -0400
Subject: [PATCH 06/26] Block the test for pl-device to check the actual
 coverage of the implementation

---
 .github/workflows/tests_lkcpu_python.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/tests_lkcpu_python.yml b/.github/workflows/tests_lkcpu_python.yml
index 0e211f2e9..83c4193e4 100644
--- a/.github/workflows/tests_lkcpu_python.yml
+++ b/.github/workflows/tests_lkcpu_python.yml
@@ -249,8 +249,8 @@ jobs:
             PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS --splits 7 --group ${{ matrix.group }} \
               --store-durations --durations-path='.github/workflows/python_lightning_kokkos_test_durations.json' --splitting-algorithm=least_duration
             mv .github/workflows/python_lightning_kokkos_test_durations.json ${{ github.workspace }}/.test_durations-${{ matrix.exec_model }}-${{ matrix.group }}
-          pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
-          pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
+          : # pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
+          : # pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv .coverage .coverage-${{ github.job }}-${{ matrix.pl_backend }}-${{ matrix.group }}
 
       - name: Upload test durations

From d186d3a3edce019ed474d45856996300dc604dd7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 10:49:46 -0400
Subject: [PATCH 07/26] Temporal fix for errors in CI

---
 .github/workflows/tests_lkcuda_python.yml | 7 ++++---
 .github/workflows/wheel_linux_aarch64.yml | 3 ++-
 .github/workflows/wheel_linux_x86_64.yml  | 3 ++-
 .github/workflows/wheel_macos_arm64.yml   | 3 ++-
 .github/workflows/wheel_macos_x86_64.yml  | 3 ++-
 5 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/tests_lkcuda_python.yml b/.github/workflows/tests_lkcuda_python.yml
index 45b012175..333fd1fc3 100644
--- a/.github/workflows/tests_lkcuda_python.yml
+++ b/.github/workflows/tests_lkcuda_python.yml
@@ -254,8 +254,8 @@ jobs:
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
           PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
-          pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
-          pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
+          # pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
+          # pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
 
       - name: Install all backend devices
@@ -275,7 +275,8 @@ jobs:
           OMP_PROC_BIND: false
         run: |
           cd main/
-          for device in lightning.qubit lightning.kokkos; do
+          # for device in lightning.qubit lightning.kokkos; do
+          for device in lightning.qubit; do
             pl-device-test --device ${device} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
             pl-device-test --device ${device} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           done
diff --git a/.github/workflows/wheel_linux_aarch64.yml b/.github/workflows/wheel_linux_aarch64.yml
index ec090d9f7..9b938a6d6 100644
--- a/.github/workflows/wheel_linux_aarch64.yml
+++ b/.github/workflows/wheel_linux_aarch64.yml
@@ -92,7 +92,8 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         arch: [aarch64]
-        pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        : # pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        pl_backend: ["lightning_qubit"]
         cibw_build: ${{ fromJson(needs.set_wheel_build_matrix.outputs.python_version) }}
         exec_model: ${{ fromJson(needs.set_wheel_build_matrix.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.set_wheel_build_matrix.outputs.kokkos_version) }}
diff --git a/.github/workflows/wheel_linux_x86_64.yml b/.github/workflows/wheel_linux_x86_64.yml
index f91ff34e8..c3b066146 100644
--- a/.github/workflows/wheel_linux_x86_64.yml
+++ b/.github/workflows/wheel_linux_x86_64.yml
@@ -100,7 +100,8 @@ jobs:
       fail-fast: false
       matrix:
         arch: [x86_64]
-        pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        : # pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        pl_backend: ["lightning_qubit"]
         cibw_build: ${{ fromJson(needs.set_wheel_build_matrix.outputs.python_version) }}
         exec_model: ${{ fromJson(needs.set_wheel_build_matrix.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.set_wheel_build_matrix.outputs.kokkos_version) }}
diff --git a/.github/workflows/wheel_macos_arm64.yml b/.github/workflows/wheel_macos_arm64.yml
index 5b1b20d3d..7c30b0eb1 100644
--- a/.github/workflows/wheel_macos_arm64.yml
+++ b/.github/workflows/wheel_macos_arm64.yml
@@ -60,7 +60,8 @@ jobs:
       matrix:
         os: [macos-12]
         arch: [arm64]
-        pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        : # pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        pl_backend: ["lightning_qubit"]
         cibw_build: ${{fromJson(needs.mac-set-matrix-arm.outputs.python_version)}}
     timeout-minutes: 30
     name: macos-latest::arm64 - ${{ matrix.pl_backend }} (Python ${{ fromJson('{ "cp39-*":"3.9","cp310-*":"3.10","cp311-*":"3.11","cp312-*":"3.12" }')[matrix.cibw_build] }})
diff --git a/.github/workflows/wheel_macos_x86_64.yml b/.github/workflows/wheel_macos_x86_64.yml
index f3aa0d900..a1e6b2395 100644
--- a/.github/workflows/wheel_macos_x86_64.yml
+++ b/.github/workflows/wheel_macos_x86_64.yml
@@ -95,7 +95,8 @@ jobs:
       matrix:
         os: [macos-12]
         arch: [x86_64]
-        pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        : # pl_backend: ["lightning_kokkos", "lightning_qubit"]
+        pl_backend: ["lightning_qubit"]
         cibw_build: ${{fromJson(needs.set_wheel_build_matrix.outputs.python_version)}}
         exec_model: ${{ fromJson(needs.set_wheel_build_matrix.outputs.exec_model) }}
         kokkos_version: ${{ fromJson(needs.set_wheel_build_matrix.outputs.kokkos_version) }}

From 7ef541653298b57a3414a0e9d52cad391d0ad30a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 11:14:55 -0400
Subject: [PATCH 08/26] Added simulate test

---
 tests/lightning_qubit/test_simulate_method.py | 137 ++++++++++++++++++
 1 file changed, 137 insertions(+)
 create mode 100644 tests/lightning_qubit/test_simulate_method.py

diff --git a/tests/lightning_qubit/test_simulate_method.py b/tests/lightning_qubit/test_simulate_method.py
new file mode 100644
index 000000000..8847e17ba
--- /dev/null
+++ b/tests/lightning_qubit/test_simulate_method.py
@@ -0,0 +1,137 @@
+# Copyright 2018-2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import math
+from typing import Sequence
+
+import numpy as np
+import pennylane as qml
+import pytest
+from conftest import PHI, THETA, LightningDevice, device_name  # tested device
+from flaky import flaky
+from pennylane.devices import DefaultExecutionConfig, DefaultQubit
+from pennylane.measurements import VarianceMP
+from scipy.sparse import csr_matrix, random_array
+
+if device_name == "lightning.qubit":
+    from pennylane_lightning.lightning_qubit._state_vector import LightningStateVector
+    from pennylane_lightning.lightning_qubit.lightning_qubit import simulate
+
+
+if device_name == "lightning.kokkos":
+    from pennylane_lightning.lightning_kokkos._state_vector import LightningKokkosStateVector as LightningStateVector
+    from pennylane_lightning.lightning_kokkos.lightning_kokkos import simulate
+
+
+if device_name != "lightning.qubit" and device_name != "lightning.kokkos":
+    pytest.skip(
+        "Exclusive tests for lightning.qubit and lightning.kokkos. Skipping.",
+        allow_module_level=True,
+    )
+
+if not LightningDevice._CPP_BINARY_AVAILABLE:
+    pytest.skip("No binary module found. Skipping.", allow_module_level=True)
+
+
+# General LightningStateVector fixture, for any number of wires.
+@pytest.fixture(
+    scope="module",
+    params=[np.complex64, np.complex128],
+)
+def lightning_sv(request):
+    def _statevector(num_wires):
+        return LightningStateVector(num_wires=num_wires, dtype=request.param)
+
+    return _statevector
+
+
+class TestSimulate:
+    """Tests for the simulate method."""
+
+    @staticmethod
+    def calculate_reference(tape):
+        dev = DefaultQubit(max_workers=1)
+        program, _ = dev.preprocess()
+        tapes, transf_fn = program([tape])
+        results = dev.execute(tapes)
+        return transf_fn(results)
+
+    def test_simple_circuit(self, lightning_sv, tol):
+        """Tests the simulate method for a simple circuit."""
+        tape = qml.tape.QuantumScript(
+            [qml.RX(THETA[0], wires=0), qml.RY(PHI[0], wires=1)],
+            [qml.expval(qml.PauliX(0))],
+            shots=None,
+        )
+        statevector = lightning_sv(num_wires=2)
+        result = simulate(circuit=tape, state=statevector)
+        reference = self.calculate_reference(tape)
+
+        assert np.allclose(result, reference, tol)
+
+    test_data_no_parameters = [
+        (100, qml.PauliZ(wires=[0]), 100),
+        (110, qml.PauliZ(wires=[1]), 110),
+        (120, qml.PauliX(0) @ qml.PauliZ(1), 120),
+    ]
+
+    @pytest.mark.parametrize("num_shots,operation,shape", test_data_no_parameters)
+    def test_sample_dimensions(self, lightning_sv, num_shots, operation, shape):
+        """Tests if the samples returned by simulate have the correct dimensions"""
+        ops = [qml.RX(1.5708, wires=[0]), qml.RX(1.5708, wires=[1])]
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=operation)], shots=num_shots)
+
+        statevector = lightning_sv(num_wires=2)
+        result = simulate(circuit=tape, state=statevector)
+
+        assert np.array_equal(result.shape, (shape,))
+
+    def test_sample_values(self, lightning_sv, tol):
+        """Tests if the samples returned by simulate have the correct values"""
+        ops = [qml.RX(1.5708, wires=[0])]
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=qml.PauliZ(0))], shots=1000)
+
+        statevector = lightning_sv(num_wires=1)
+
+        result = simulate(circuit=tape, state=statevector)
+
+        assert np.allclose(result**2, 1, atol=tol, rtol=0)
+
+    @pytest.mark.skipif(
+        device_name == "lightning.kokkos",
+        reason=f"Device {device_name} does not have an mcmc option.",
+    )
+    @pytest.mark.parametrize("mcmc", [True, False])
+    @pytest.mark.parametrize("kernel", ["Local", "NonZeroRandom"])
+    def test_sample_values_with_mcmc(self, lightning_sv, tol, mcmc, kernel):
+        """Tests if the samples returned by simulate have the correct values"""
+        ops = [qml.RX(1.5708, wires=[0])]
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=qml.PauliZ(0))], shots=1000)
+
+        statevector = lightning_sv(num_wires=1)
+
+        mcmc_param = {
+            "mcmc": mcmc,
+            "kernel_name": kernel,
+            "num_burnin": 100,
+        }
+
+        execution_config = DefaultExecutionConfig
+
+        result = simulate(
+            circuit=tape, state=statevector, mcmc=mcmc_param, postselect_mode=execution_config
+        )
+
+        assert np.allclose(result**2, 1, atol=tol, rtol=0)

From d7f06103489f6a84db0c3b0e4a578a13feb14945 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 11:17:28 -0400
Subject: [PATCH 09/26] apply format

---
 tests/lightning_qubit/test_simulate_method.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/lightning_qubit/test_simulate_method.py b/tests/lightning_qubit/test_simulate_method.py
index 8847e17ba..75f5c5da6 100644
--- a/tests/lightning_qubit/test_simulate_method.py
+++ b/tests/lightning_qubit/test_simulate_method.py
@@ -31,7 +31,9 @@
 
 
 if device_name == "lightning.kokkos":
-    from pennylane_lightning.lightning_kokkos._state_vector import LightningKokkosStateVector as LightningStateVector
+    from pennylane_lightning.lightning_kokkos._state_vector import (
+        LightningKokkosStateVector as LightningStateVector,
+    )
     from pennylane_lightning.lightning_kokkos.lightning_kokkos import simulate
 
 

From 2a31aafac2f12591d407723de40278f041fdc246 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 14:21:25 -0400
Subject: [PATCH 10/26] Update tests/lightning_qubit/test_simulate_method.py

Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
---
 tests/lightning_qubit/test_simulate_method.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/lightning_qubit/test_simulate_method.py b/tests/lightning_qubit/test_simulate_method.py
index 75f5c5da6..8fffc721e 100644
--- a/tests/lightning_qubit/test_simulate_method.py
+++ b/tests/lightning_qubit/test_simulate_method.py
@@ -112,7 +112,7 @@ def test_sample_values(self, lightning_sv, tol):
         assert np.allclose(result**2, 1, atol=tol, rtol=0)
 
     @pytest.mark.skipif(
-        device_name == "lightning.kokkos",
+        device_name != "lightning.qubit",
         reason=f"Device {device_name} does not have an mcmc option.",
     )
     @pytest.mark.parametrize("mcmc", [True, False])

From 72faa2fa679ff1d22c4d5ad81c222f370a37d9f9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 14:43:13 -0400
Subject: [PATCH 11/26] Include the changes introduced in #814 for measurements

---
 .../lightning_kokkos/_measurements.py         | 44 +++++--------------
 1 file changed, 10 insertions(+), 34 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py
index f641c0fdd..d505c8883 100644
--- a/pennylane_lightning/lightning_kokkos/_measurements.py
+++ b/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,

From 5226b8ca3a44c968028be67377cc427f45c98964 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 15:11:58 -0400
Subject: [PATCH 12/26] remove _asarray method

---
 .../lightning_kokkos/_state_vector.py         | 21 -------------------
 1 file changed, 21 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index b203b3fd7..3a97f9a16 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -135,30 +135,9 @@ def state(self):
         [0.+0.j 1.+0.j]
         """
         state = np.zeros(2**self._num_wires, dtype=self.dtype)
-        state = self._asarray(state, dtype=self.dtype)
         self.sync_d2h(state)
         return state
 
-    @staticmethod
-    def _asarray(arr, dtype=None):
-        arr = np.asarray(arr)  # arr is not copied
-
-        if arr.dtype.kind not in ["f", "c"]:
-            return arr
-
-        if not dtype:
-            dtype = arr.dtype
-
-        # We allocate a new aligned memory and copy data to there if alignment
-        # or dtype mismatches
-        # Note that get_alignment does not necessarily return CPUMemoryModel(Unaligned) even for
-        # numpy allocated memory as the memory location happens to be aligned.
-        if arr.dtype != dtype:
-            new_arr = allocate_aligned_array(arr.size, np.dtype(dtype), False).reshape(arr.shape)
-            np.copyto(new_arr, arr)
-            arr = new_arr
-        return arr
-
     def sync_h2d(self, state_vector):
         """Copy the state vector data on host provided by the user to the state
         vector on the device

From ca348338fe8396ad80660fc67fc4ccb439180ae7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Fri, 26 Jul 2024 15:34:52 -0400
Subject: [PATCH 13/26] trigger CIs


From b60a75147b3307b12c2e55605c54a547f8735f84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 11:01:39 -0400
Subject: [PATCH 14/26] Simplify apply_operations method

Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
---
 pennylane_lightning/lightning_kokkos/_state_vector.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index 3a97f9a16..1b7c70836 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -446,10 +446,9 @@ def apply_operations(
         if operations:  # make sure operations[0] exists
             if isinstance(operations[0], StatePrep):
                 self._apply_state_vector(operations[0].parameters[0].copy(), operations[0].wires)
-                operations = operations[1:]
             elif isinstance(operations[0], BasisState):
                 self._apply_basis_state(operations[0].parameters[0], operations[0].wires)
-                operations = operations[1:]
+            operations = operations[1:]
 
         self._apply_lightning(
             operations, mid_measurements=mid_measurements, postselect_mode=postselect_mode

From c1fa3eb13446cb49d54fc7747e124a42461d57fe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 15:43:44 -0400
Subject: [PATCH 15/26] Apply review comments

---
 .../lightning_kokkos/_measurements.py         |  2 +-
 .../lightning_kokkos/_state_vector.py         | 61 ++++---------------
 .../test_measurements_class.py                |  5 +-
 3 files changed, 16 insertions(+), 52 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py
index d505c8883..532c2248e 100644
--- a/pennylane_lightning/lightning_kokkos/_measurements.py
+++ b/pennylane_lightning/lightning_kokkos/_measurements.py
@@ -66,7 +66,7 @@ def __init__(
 
     @property
     def qubit_state(self):
-        """Returns a handle to the LightningStateVector class."""
+        """Returns a handle to the LightningStateVector object."""
         return self._qubit_state
 
     @property
diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index 1b7c70836..ad11f284a 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -86,9 +86,8 @@ def __init__(
         elif isinstance(kokkos_args, InitializationSettings):
             self._kokkos_state = self._state_dtype()(self.num_wires, kokkos_args)
         else:
-            type0 = type(InitializationSettings())
             raise TypeError(
-                f"Argument kokkos_args must be of type {type0} but it is of {type(kokkos_args)}."
+                f"Argument kokkos_args must be of type {type(InitializationSettings())} but it is of {type(kokkos_args)}."
             )
         self._sync = sync
 
@@ -192,14 +191,6 @@ def _state_dtype(self):
         """
         return StateVectorC128 if self.dtype == np.complex128 else StateVectorC64
 
-    def _create_basis_state(self, index):
-        """Return a computational basis state over all wires.
-
-        Args:
-            index (int): integer representing the computational basis state.
-        """
-        self._kokkos_state.setBasisState(index)
-
     def reset_state(self):
         """Reset the device's state"""
         # init the state vector to |00..0>
@@ -298,7 +289,8 @@ def _apply_basis_state(self, state, wires):
         Note: This function does not support broadcasted inputs yet.
         """
         num = self._get_basis_state_index(state, wires)
-        self._create_basis_state(num)
+        # Return a computational basis state over all wires.
+        self._kokkos_state.setBasisState(num)
 
     def _apply_lightning_controlled(self, operation):
         """Apply an arbitrary controlled operation to the state tensor.
@@ -311,46 +303,15 @@ def _apply_lightning_controlled(self, operation):
         """
         state = self.state_vector
 
-        basename = operation.base.name
-        method = getattr(state, f"{basename}", None)
         control_wires = list(operation.control_wires)
         control_values = operation.control_values
-        #  ----------------------------------------
-        # Original:
-        # target_wires = list(operation.target_wires)
-
-        # Specific for Kokkos:
         name = operation.name
-        #  ----------------------------------------
-        if method is not None:  # apply n-controlled specialized gate
-            inv = False
-            #  ----------------------------------------
-            # Original:
-            # param = operation.parameters
-            # method(control_wires, control_values, target_wires, inv, param)
-
-            # Specific for Kokkos:
-            param = operation.parameters[0]
-            wires = self.wires.indices(operation.wires)
-            matrix = global_phase_diagonal(param, self.wires, control_wires, control_values)
-            state.apply(name, wires, inv, [[param]], matrix)
-            #  ----------------------------------------
-
-        else:  # apply gate as an n-controlled matrix
-            #  ----------------------------------------
-            # Original:
-            # method = getattr(state, "applyControlledMatrix")
-            # method(
-            #     qml.matrix(operation.base),
-            #     control_wires,
-            #     control_values,
-            #     target_wires,
-            #     False,
-            # )
-
-            # Specific for Kokkos:
-            raise ValueError("Unsupported apply Controlled Matrix")
-            #  ----------------------------------------
+        # Apply GlobalPhase
+        inv = False
+        param = operation.parameters[0]
+        wires = self.wires.indices(operation.wires)
+        matrix = global_phase_diagonal(param, self.wires, control_wires, control_values)
+        state.apply(name, wires, inv, [[param]], matrix)
 
     def _apply_lightning_midmeasure(
         self, operation: MidMeasureMP, mid_measurements: dict, postselect_mode: str
@@ -424,9 +385,9 @@ def _apply_lightning(
             elif (
                 isinstance(operation, qml.ops.Controlled)
                 and isinstance(operation.base, qml.GlobalPhase)
-                # Specific for Kokkos:
-                # Kokkos do not support the controlled gates except for GlobalPhase
             ):  # apply n-controlled gate
+
+                # Kokkos do not support the controlled gates except for GlobalPhase
                 self._apply_lightning_controlled(operation)
             else:  # apply gate as a matrix
                 # Inverse can be set to False since qml.matrix(operation) is already in
diff --git a/tests/lightning_qubit/test_measurements_class.py b/tests/lightning_qubit/test_measurements_class.py
index 33aad93f9..a4fe46857 100644
--- a/tests/lightning_qubit/test_measurements_class.py
+++ b/tests/lightning_qubit/test_measurements_class.py
@@ -487,7 +487,7 @@ def test_single_return_value(self, measurement, observable, lightning_sv, tol):
         assert np.allclose(result, expected, max(tol, 1.0e-4))
 
     @flaky(max_runs=5)
-    @pytest.mark.parametrize("shots", [None, 1000000])
+    @pytest.mark.parametrize("shots", [None, 1000000,(900000,900000)])
     @pytest.mark.parametrize("measurement", [qml.expval, qml.probs, qml.var])
     @pytest.mark.parametrize(
         "obs0_",
@@ -582,6 +582,9 @@ def test_double_return_value(self, shots, measurement, obs0_, obs1_, lightning_s
         # a few tests may fail in single precision, and hence we increase the tolerance
         dtol = tol if shots is None else max(tol, 1.0e-2)
         for r, e in zip(result, expected):
+            if isinstance(shots,tuple) and  isinstance(r[0], np.ndarray):
+                r = np.concatenate(r)
+                e = np.concatenate(e)
             assert np.allclose(r, e, atol=dtol, rtol=dtol)
 
     @pytest.mark.parametrize(

From 941377bae1722c321511de13e8ffafc73aa7e661 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 15:46:19 -0400
Subject: [PATCH 16/26] apply format

---
 pennylane_lightning/lightning_kokkos/_state_vector.py | 5 ++---
 tests/lightning_qubit/test_measurements_class.py      | 4 ++--
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index ad11f284a..4196d8be8 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -382,9 +382,8 @@ def _apply_lightning(
             elif method is not None:  # apply specialized gate
                 param = operation.parameters
                 method(wires, invert_param, param)
-            elif (
-                isinstance(operation, qml.ops.Controlled)
-                and isinstance(operation.base, qml.GlobalPhase)
+            elif isinstance(operation, qml.ops.Controlled) and isinstance(
+                operation.base, qml.GlobalPhase
             ):  # apply n-controlled gate
 
                 # Kokkos do not support the controlled gates except for GlobalPhase
diff --git a/tests/lightning_qubit/test_measurements_class.py b/tests/lightning_qubit/test_measurements_class.py
index a4fe46857..afbaa6dbe 100644
--- a/tests/lightning_qubit/test_measurements_class.py
+++ b/tests/lightning_qubit/test_measurements_class.py
@@ -487,7 +487,7 @@ def test_single_return_value(self, measurement, observable, lightning_sv, tol):
         assert np.allclose(result, expected, max(tol, 1.0e-4))
 
     @flaky(max_runs=5)
-    @pytest.mark.parametrize("shots", [None, 1000000,(900000,900000)])
+    @pytest.mark.parametrize("shots", [None, 1000000, (900000, 900000)])
     @pytest.mark.parametrize("measurement", [qml.expval, qml.probs, qml.var])
     @pytest.mark.parametrize(
         "obs0_",
@@ -582,7 +582,7 @@ def test_double_return_value(self, shots, measurement, obs0_, obs1_, lightning_s
         # a few tests may fail in single precision, and hence we increase the tolerance
         dtol = tol if shots is None else max(tol, 1.0e-2)
         for r, e in zip(result, expected):
-            if isinstance(shots,tuple) and  isinstance(r[0], np.ndarray):
+            if isinstance(shots, tuple) and isinstance(r[0], np.ndarray):
                 r = np.concatenate(r)
                 e = np.concatenate(e)
             assert np.allclose(r, e, atol=dtol, rtol=dtol)

From aef2f9777800cb34e9f81e50fa2690ca09b01cc7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 16:19:05 -0400
Subject: [PATCH 17/26] Fix bug

---
 pennylane_lightning/lightning_kokkos/_state_vector.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index 4196d8be8..ba7260625 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -406,9 +406,10 @@ def apply_operations(
         if operations:  # make sure operations[0] exists
             if isinstance(operations[0], StatePrep):
                 self._apply_state_vector(operations[0].parameters[0].copy(), operations[0].wires)
+                operations = operations[1:]
             elif isinstance(operations[0], BasisState):
                 self._apply_basis_state(operations[0].parameters[0], operations[0].wires)
-            operations = operations[1:]
+                operations = operations[1:]
 
         self._apply_lightning(
             operations, mid_measurements=mid_measurements, postselect_mode=postselect_mode

From 3134460e63b82a07563eae5ef6e8296d7fd298c5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 16:44:44 -0400
Subject: [PATCH 18/26] correct comment format for
 github/workflows/tests_lkcuda_python.yml

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
---
 .github/workflows/tests_lkcuda_python.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/tests_lkcuda_python.yml b/.github/workflows/tests_lkcuda_python.yml
index 333fd1fc3..2f5795ba5 100644
--- a/.github/workflows/tests_lkcuda_python.yml
+++ b/.github/workflows/tests_lkcuda_python.yml
@@ -254,8 +254,8 @@ jobs:
           cd main/
           DEVICENAME=`echo ${{ matrix.pl_backend }} | sed "s/_/./g"`
           PL_DEVICE=${DEVICENAME} python -m pytest tests/ -k "not test_native_mcm" $COVERAGE_FLAGS
-          # pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
-          # pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
+          : # pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
+          : # pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
           mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
 
       - name: Install all backend devices

From e1f09454505e9da3f786e2bff693d7a8ca7f5dbb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 17:05:12 -0400
Subject: [PATCH 19/26] Apply suggestions from code review in doc for d2h in
 kokkos

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
---
 pennylane_lightning/lightning_kokkos/_state_vector.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index ba7260625..cc9c71f95 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -164,7 +164,7 @@ def sync_d2h(self, state_vector):
         by the user
 
         Args:
-            state_vector(array[complex]): the state vector array on host
+            state_vector(array[complex]): the state vector array on device
 
 
         **Example**

From 8636a2b33960667b6db0512603efbce753b9ffa4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 17:06:39 -0400
Subject: [PATCH 20/26] change the doc for kokkos_configuration

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
---
 pennylane_lightning/lightning_kokkos/_state_vector.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index cc9c71f95..724fb2ed4 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -179,8 +179,9 @@ def sync_d2h(self, state_vector):
         self._kokkos_state.DeviceToHost(state_vector.ravel(order="C"))
 
     def _kokkos_configuration(self):
-        """Set the default configuration of the kokkos device.
-        Returns: kokkos configuration
+        """Get the default configuration of the kokkos device.
+
+        Returns: The `lightning.kokkos` device configuration
         """
         return print_configuration()
 

From d64f50b725389de41905f00d195dbb705ef1fc66 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 17:08:46 -0400
Subject: [PATCH 21/26] avoid double call of reset_state

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
---
 pennylane_lightning/lightning_kokkos/lightning_kokkos.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
index 5b21fd3ed..8917566c6 100644
--- a/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
+++ b/pennylane_lightning/lightning_kokkos/lightning_kokkos.py
@@ -65,7 +65,6 @@ def simulate(
 
     Note that this function can return measurements for non-commuting observables simultaneously.
     """
-    state.reset_state()
     has_mcm = any(isinstance(op, MidMeasureMP) for op in circuit.operations)
     if circuit.shots and has_mcm:
         results = []
@@ -87,7 +86,7 @@ def simulate(
                 )
             )
         return tuple(results)
-
+    state.reset_state()
     final_state = state.get_final_state(circuit)
     return LightningKokkosMeasurements(final_state).measure_final_state(circuit)
 

From fdab75f8df0b797a168255baa3b598dbe16a1437 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 17:11:51 -0400
Subject: [PATCH 22/26] reorder the import modules sequence

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
---
 .../lightning_kokkos/_state_vector.py         | 22 +++++++++----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index 724fb2ed4..157299cde 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -15,17 +15,6 @@
 Class implementation for lightning_kokkos state-vector manipulation.
 """
 
-from itertools import product
-
-import numpy as np
-import pennylane as qml
-from pennylane import BasisState, DeviceError, StatePrep
-from pennylane.measurements import MidMeasureMP
-from pennylane.ops import Conditional
-from pennylane.ops.op_math import Adjoint
-from pennylane.tape import QuantumScript
-from pennylane.wires import Wires
-
 try:
     from pennylane_lightning.lightning_kokkos_ops import (
         InitializationSettings,
@@ -37,6 +26,17 @@
 except ImportError:
     pass  # Should be a complaint when kokkos_ops module is not available.
 
+from itertools import product
+
+import numpy as np
+import pennylane as qml
+from pennylane import BasisState, DeviceError, StatePrep
+from pennylane.measurements import MidMeasureMP
+from pennylane.ops import Conditional
+from pennylane.ops.op_math import Adjoint
+from pennylane.tape import QuantumScript
+from pennylane.wires import Wires
+
 from pennylane_lightning.core._serialize import global_phase_diagonal
 
 from ._measurements import LightningKokkosMeasurements

From c836bacacc15f89b7aeb2314d5b099e33482d8dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 17:26:58 -0400
Subject: [PATCH 23/26] regroup attributes in state_vector

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>
---
 pennylane_lightning/lightning_kokkos/_state_vector.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane_lightning/lightning_kokkos/_state_vector.py b/pennylane_lightning/lightning_kokkos/_state_vector.py
index 157299cde..d1fd75a48 100644
--- a/pennylane_lightning/lightning_kokkos/_state_vector.py
+++ b/pennylane_lightning/lightning_kokkos/_state_vector.py
@@ -71,6 +71,7 @@ def __init__(
         self._dtype = dtype
 
         self._kokkos_config = {}
+        self._sync = sync
 
         if dtype not in [np.complex64, np.complex128]:  # pragma: no cover
             raise TypeError(f"Unsupported complex type: {dtype}")
@@ -89,7 +90,6 @@ def __init__(
             raise TypeError(
                 f"Argument kokkos_args must be of type {type(InitializationSettings())} but it is of {type(kokkos_args)}."
             )
-        self._sync = sync
 
         if not self._kokkos_config:
             self._kokkos_config = self._kokkos_configuration()

From bccc56bb6f96b3cf2a7e3e905c2e666424c27ffb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 17:28:56 -0400
Subject: [PATCH 24/26] Change doc for Lightning Kokkos Measurements

---
 .../lightning_kokkos/_measurements.py                | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/pennylane_lightning/lightning_kokkos/_measurements.py b/pennylane_lightning/lightning_kokkos/_measurements.py
index 532c2248e..23dc3aa70 100644
--- a/pennylane_lightning/lightning_kokkos/_measurements.py
+++ b/pennylane_lightning/lightning_kokkos/_measurements.py
@@ -21,7 +21,6 @@
 except ImportError:
     pass
 
-from functools import reduce
 from typing import Callable, List, Union
 
 import numpy as np
@@ -48,12 +47,12 @@
 
 
 class LightningKokkosMeasurements:
-    """Lightning Measurements class
+    """Lightning Kokkos Measurements class
 
-    Measures the state provided by the LightningStateVector class.
+    Measures the state provided by the LightningKokkosStateVector class.
 
     Args:
-        qubit_state(LightningStateVector): Lightning state-vector class containing the state vector to be measured.
+        qubit_state(LightningKokkosStateVector): Lightning state-vector class containing the state vector to be measured.
     """
 
     def __init__(
@@ -66,7 +65,7 @@ def __init__(
 
     @property
     def qubit_state(self):
-        """Returns a handle to the LightningStateVector object."""
+        """Returns a handle to the LightningKokkosStateVector object."""
         return self._qubit_state
 
     @property
@@ -75,7 +74,7 @@ def dtype(self):
         return self._dtype
 
     def _measurement_dtype(self):
-        """Binding to Lightning Measurements C++ class.
+        """Binding to Lightning Kokkos Measurements C++ class.
 
         Returns: the Measurements class
         """
@@ -265,7 +264,6 @@ def measure_final_state(self, circuit: QuantumScript, mid_measurements=None) ->
 
         return results
 
-    # pylint:disable = too-many-arguments
     def measure_with_samples(
         self,
         measurements: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]],

From 9a094ca97ccb7d6b1c350034e7ad4516e2a6d97d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Mon, 29 Jul 2024 18:09:42 -0400
Subject: [PATCH 25/26] trigger CIs


From c4860a931f0f9fad1faaf3230c34275c7db159e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Alfredo=20Nu=C3=B1ez=20Meneses?=
 <alfredo.nunez@xanadu.ai>
Date: Tue, 30 Jul 2024 09:04:02 -0400
Subject: [PATCH 26/26] remove useless lines from test_state_vector

Co-authored-by: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
---
 tests/lightning_qubit/test_state_vector_class.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/lightning_qubit/test_state_vector_class.py b/tests/lightning_qubit/test_state_vector_class.py
index 5cea7c3e0..f859a0022 100644
--- a/tests/lightning_qubit/test_state_vector_class.py
+++ b/tests/lightning_qubit/test_state_vector_class.py
@@ -45,9 +45,7 @@
 
 @pytest.mark.parametrize("num_wires", range(4))
 @pytest.mark.parametrize("dtype", [np.complex64, np.complex128])
-# @pytest.mark.parametrize("device_name", ["lightning.qubit"])
-@pytest.mark.parametrize("device_name", [device_name])
-def test_device_name_and_init(num_wires, dtype, device_name):
+def test_device_name_and_init(num_wires, dtype):
     """Test the class initialization and returned properties."""
     state_vector = LightningStateVector(num_wires, dtype=dtype, device_name=device_name)
     assert state_vector.dtype == dtype