Allow lightning.qubit/kokkos::generate_samples to take in seeds to make the generated samples deterministic (#927)

### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to
the [`tests`](../tests) directory!

- [x] All new functions and code must be clearly commented and
documented.
If you do make documentation changes, make sure that the docs build and 
render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `.github/CHANGELOG.md` file, summarizing
the change, and including a link back to the PR.

- [x] Ensure that code is properly formatted by running `make format`. 

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.


------------------------------------------------------------------------------------------------------------

**Context:**
[A while ago](PennyLaneAI/catalyst#936) a new
`seed` option to `qjit` was added. The seed was used to make measurement
results deterministic, but samples were still probabilistic. This is
because within a `qjit` context, [measurements were controlled from the
catalyst
repo](https://github.com/PennyLaneAI/catalyst/blob/a580bada575793b780d5366aa77dff6157cd4f93/runtime/lib/backend/common/Utils.hpp#L274)
, but samples were controlled by lightning.

To resolve stochastically failing tests (i.e. flaky tests) in catalyst,
we add seeding for samples in lightning.

**Description of the Change:**
When `qjit(seed=...)` receives a (unsigned 32 bit int) seed value from
the user, the seed gets propagated through mlir and [generates a
`std::mt19937` rng instance in the catalyst execution
context](https://github.com/PennyLaneAI/catalyst/blob/934726fe750043886415953dbd89a4c4ddeb9a80/runtime/lib/capi/ExecutionContext.hpp#L268).
This rng instance eventually becomes a field of the
`Catalyst::Runtime::Simulator::LightningSimulator` (and kokkos) class
[catalyst/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.hpp](https://github.com/PennyLaneAI/catalyst/blob/a580bada575793b780d5366aa77dff6157cd4f93/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.hpp#L54).

To seed samples, catalyst uses this device rng instance on the state
vector's `generate_samples` methods:
PennyLaneAI/catalyst#1164.

In lightning, the `generate_samples` method now takes in a seeding
number. The catalyst devices pass in a seed into the lightning
`generate_samples`; this seed is created deterministically from the
aforementioned already seeded catalyst context rng instance. This makes
the generated samples deterministc.


**Benefits:** 
Fewer (hopefully no) stochatically failing frontend tests in catalyst.

**Possible Drawbacks:**

**Related GitHub Issues:**
[sc-72878]

---------

Co-authored-by: ringo-but-quantum <[email protected]>

Author: paul0403

.github/CHANGELOG.md

@@ -40,6 +40,9 @@
 
 ### Improvements
 
+* The `generate_samples` methods of lightning.{qubit/kokkos} can now take in a seed number to make the generated samples deterministic. This can be useful when, among other things, fixing flaky tests in CI.
+  [(#927)](https://github.com/PennyLaneAI/pennylane-lightning/pull/927)
+
 * Always decompose `qml.QFT` in Lightning.
   [(#924)](https://github.com/PennyLaneAI/pennylane-lightning/pull/924)
 
@@ -111,7 +114,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Ali Asadi, Amintor Dusko, Luis Alfredo Nuñez Meneses, Vincent Michaud-Rioux, Lee J. O'Riordan, Mudit Pandey, Shuli Shu
+Ali Asadi, Amintor Dusko, Luis Alfredo Nuñez Meneses, Vincent Michaud-Rioux, Lee J. O'Riordan, Mudit Pandey, Shuli Shu, Haochen Paul Wang
 
 ---
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.39.0-dev37"
+__version__ = "0.39.0-dev38"

pennylane_lightning/core/src/measurements/MeasurementsBase.hpp

@@ -77,7 +77,6 @@ template <class StateVectorT, class Derived> class MeasurementsBase {
     /**
      * @brief Randomly set the seed of the internal random generator
      *
-     * @param seed Seed
      */
     void setRandomSeed() {
         std::random_device rd;

pennylane_lightning/core/src/measurements/tests/Test_MeasurementsBase.cpp

@@ -20,6 +20,7 @@ using Pennylane::Util::isApproxEqual;
 } // namespace
 /// @endcond
 #include <algorithm>
+#include <optional>
 #include <string>
 
 #ifdef _ENABLE_PLQUBIT
@@ -1251,7 +1252,9 @@ TEST_CASE("Var Shot- TensorProdObs", "[MeasurementsBase][Observables]") {
         testTensorProdObsVarShot<TestStateVectorBackends>();
     }
 }
-template <typename TypeList> void testSamples() {
+
+template <typename TypeList>
+void testSamples(const std::optional<std::size_t> &seed = std::nullopt) {
     if constexpr (!std::is_same_v<TypeList, void>) {
         using StateVectorT = typename TypeList::Type;
         using PrecisionT = typename StateVectorT::PrecisionT;
@@ -1281,7 +1284,10 @@ template <typename TypeList> void testSamples() {
         std::size_t num_qubits = 3;
         std::size_t N = std::pow(2, num_qubits);
         std::size_t num_samples = 100000;
-        auto &&samples = Measurer.generate_samples(num_samples);
+        auto &&samples =
+            seed.has_value()
+                ? Measurer.generate_samples(num_samples, seed.value())
+                : Measurer.generate_samples(num_samples);
 
         std::vector<std::size_t> counts(N, 0);
         std::vector<std::size_t> samples_decimal(num_samples, 0);
@@ -1307,7 +1313,7 @@ template <typename TypeList> void testSamples() {
             REQUIRE_THAT(probabilities,
                          Catch::Approx(expected_probabilities).margin(.05));
         }
-        testSamples<typename TypeList::Next>();
+        testSamples<typename TypeList::Next>(seed);
     }
 }
 
@@ -1317,6 +1323,12 @@ TEST_CASE("Samples", "[MeasurementsBase]") {
     }
 }
 
+TEST_CASE("Seeded samples", "[MeasurementsBase]") {
+    if constexpr (BACKEND_FOUND) {
+        testSamples<TestStateVectorBackends>(37);
+    }
+}
+
 template <typename TypeList> void testSamplesCountsObs() {
     if constexpr (!std::is_same_v<TypeList, void>) {
         using StateVectorT = typename TypeList::Type;
@@ -1729,4 +1741,4 @@ TEST_CASE("Measure Shot - SparseHObs ", "[MeasurementsBase][Observables]") {
     if constexpr (BACKEND_FOUND) {
         testSparseHObsMeasureShot<TestStateVectorBackends>();
     }
-}
+}

pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPU.hpp

@@ -25,6 +25,7 @@
 #include <cuda.h>
 #include <cusparse.h>
 #include <custatevec.h> // custatevecApplyMatrix
+#include <optional>
 #include <random>
 #include <type_traits>
 #include <unordered_map>
@@ -218,7 +219,9 @@ class Measurements final
      * be accessed using the stride sample_id*num_qubits, where sample_id is a
      * number between 0 and num_samples-1.
      */
-    auto generate_samples(std::size_t num_samples) -> std::vector<std::size_t> {
+    auto generate_samples(std::size_t num_samples,
+                          const std::optional<std::size_t> &seed = std::nullopt)
+        -> std::vector<std::size_t> {
         std::vector<double> rand_nums(num_samples);
         custatevecSamplerDescriptor_t sampler;
 
@@ -238,7 +241,11 @@ class Measurements final
             data_type = CUDA_C_32F;
         }
 
-        this->setRandomSeed();
+        if (seed.has_value()) {
+            this->setSeed(seed.value());
+        } else {
+            this->setRandomSeed();
+        }
         std::uniform_real_distribution<PrecisionT> dis(0.0, 1.0);
         for (std::size_t n = 0; n < num_samples; n++) {
             rand_nums[n] = dis(this->rng);

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp

@@ -342,13 +342,22 @@ void LightningKokkosSimulator::PartialProbs(
     std::move(dv_probs.begin(), dv_probs.end(), probs.begin());
 }
 
-void LightningKokkosSimulator::Sample(DataView<double, 2> &samples,
-                                      std::size_t shots) {
+std::vector<size_t> LightningKokkosSimulator::GenerateSamples(size_t shots) {
+    // generate_samples is a member function of the Measures class.
     Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
+
     // PL-Lightning-Kokkos generates samples using the alias method.
     // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+    if (this->gen) {
+        return m.generate_samples(shots, (*(this->gen))());
+    }
+    return m.generate_samples(shots);
+}
+
+void LightningKokkosSimulator::Sample(DataView<double, 2> &samples,
+                                      std::size_t shots) {
+    auto li_samples = this->GenerateSamples(shots);
 
     RT_FAIL_IF(samples.size() != li_samples.size(),
                "Invalid size for the pre-allocated samples");
@@ -381,13 +390,7 @@ void LightningKokkosSimulator::PartialSample(
     // get device wires
     auto &&dev_wires = getDeviceWires(wires);
 
-    // generate_samples is a member function of the MeasuresKokkos class.
-    Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
-        *(this->device_sv)};
-
-    // PL-Lightning-Kokkos generates samples using the alias method.
-    // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+    auto li_samples = this->GenerateSamples(shots);
 
     // The lightning samples are layed out as a single vector of size
     // shots*qubits, where each element represents a single bit. The
@@ -411,13 +414,7 @@ void LightningKokkosSimulator::Counts(DataView<double, 1> &eigvals,
     RT_FAIL_IF(eigvals.size() != numElements || counts.size() != numElements,
                "Invalid size for the pre-allocated counts");
 
-    // generate_samples is a member function of the MeasuresKokkos class.
-    Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
-        *(this->device_sv)};
-
-    // PL-Lightning-Kokkos generates samples using the alias method.
-    // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+    auto li_samples = this->GenerateSamples(shots);
 
     // Fill the eigenvalues with the integer representation of the corresponding
     // computational basis bitstring. In the future, eigenvalues can also be
@@ -455,13 +452,7 @@ void LightningKokkosSimulator::PartialCounts(
     // get device wires
     auto &&dev_wires = getDeviceWires(wires);
 
-    // generate_samples is a member function of the MeasuresKokkos class.
-    Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
-        *(this->device_sv)};
-
-    // PL-Lightning-Kokkos generates samples using the alias method.
-    // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    auto li_samples = m.generate_samples(shots);
+    auto li_samples = this->GenerateSamples(shots);
 
     // Fill the eigenvalues with the integer representation of the corresponding
     // computational basis bitstring. In the future, eigenvalues can also be

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.hpp

@@ -96,6 +96,8 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
         return res;
     }
 
+    auto GenerateSamples(size_t shots) -> std::vector<size_t>;
+
   public:
     explicit LightningKokkosSimulator(const std::string &kwargs = "{}") {
         auto &&args = Catalyst::Runtime::parse_kwargs(kwargs);

pennylane_lightning/core/src/simulators/lightning_kokkos/catalyst/tests/Test_LightningKokkosMeasures.cpp

@@ -1754,26 +1754,71 @@ TEST_CASE("Counts and PartialCounts tests with numWires=0-4 shots=100",
 }
 
 TEST_CASE("Measurement with a seeded device", "[Measures]") {
-    for (std::size_t _ = 0; _ < 5; _++) {
-        std::unique_ptr<LKSimulator> sim = std::make_unique<LKSimulator>();
-        std::unique_ptr<LKSimulator> sim1 = std::make_unique<LKSimulator>();
+    std::array<std::unique_ptr<LKSimulator>, 2> sims;
+    std::vector<std::mt19937> gens{std::mt19937{37}, std::mt19937{37}};
 
-        std::mt19937 gen(37);
-        sim->SetDevicePRNG(&gen);
+    auto circuit = [](LKSimulator &sim, std::mt19937 &gen) {
+        sim.SetDevicePRNG(&gen);
         std::vector<intptr_t> Qs;
         Qs.reserve(1);
-        Qs.push_back(sim->AllocateQubit());
-        sim->NamedOperation("Hadamard", {}, {Qs[0]}, false);
-        auto m = sim->Measure(Qs[0]);
-
-        std::mt19937 gen1(37);
-        sim1->SetDevicePRNG(&gen1);
-        std::vector<intptr_t> Qs1;
-        Qs1.reserve(1);
-        Qs1.push_back(sim1->AllocateQubit());
-        sim1->NamedOperation("Hadamard", {}, {Qs1[0]}, false);
-        auto m1 = sim1->Measure(Qs1[0]);
-
-        CHECK(*m == *m1);
+        Qs.push_back(sim.AllocateQubit());
+        sim.NamedOperation("Hadamard", {}, {Qs[0]}, false);
+        auto m = sim.Measure(Qs[0]);
+        return m;
+    };
+
+    for (std::size_t trial = 0; trial < 5; trial++) {
+        sims[0] = std::make_unique<LKSimulator>();
+        sims[1] = std::make_unique<LKSimulator>();
+
+        auto m0 = circuit(*(sims[0]), gens[0]);
+        auto m1 = circuit(*(sims[1]), gens[1]);
+
+        CHECK(*m0 == *m1);
+    }
+}
+
+TEST_CASE("Sample with a seeded device", "[Measures]") {
+    std::size_t shots = 100;
+    std::array<std::unique_ptr<LKSimulator>, 2> sims;
+    std::vector<std::vector<double>> sample_vec(2,
+                                                std::vector<double>(shots * 4));
+
+    std::vector<MemRefT<double, 2>> buffers{
+        MemRefT<double, 2>{
+            sample_vec[0].data(), sample_vec[0].data(), 0, {shots, 1}, {1, 1}},
+        MemRefT<double, 2>{
+            sample_vec[1].data(), sample_vec[1].data(), 0, {shots, 1}, {1, 1}},
+    };
+    std::vector<DataView<double, 2>> views{
+        DataView<double, 2>(buffers[0].data_aligned, buffers[0].offset,
+                            buffers[0].sizes, buffers[0].strides),
+        DataView<double, 2>(buffers[1].data_aligned, buffers[1].offset,
+                            buffers[1].sizes, buffers[1].strides)};
+
+    std::vector<std::mt19937> gens{std::mt19937{37}, std::mt19937{37}};
+
+    auto circuit = [shots](LKSimulator &sim, DataView<double, 2> &view,
+                           std::mt19937 &gen) {
+        sim.SetDevicePRNG(&gen);
+        std::vector<intptr_t> Qs;
+        Qs.reserve(1);
+        Qs.push_back(sim.AllocateQubit());
+        sim.NamedOperation("Hadamard", {}, {Qs[0]}, false);
+        sim.NamedOperation("RX", {0.5}, {Qs[0]}, false);
+        sim.Sample(view, shots);
+    };
+
+    for (std::size_t trial = 0; trial < 5; trial++) {
+        sims[0] = std::make_unique<LKSimulator>();
+        sims[1] = std::make_unique<LKSimulator>();
+
+        for (std::size_t sim_idx = 0; sim_idx < sims.size(); sim_idx++) {
+            circuit(*(sims[sim_idx]), views[sim_idx], gens[sim_idx]);
+        }
+
+        for (std::size_t i = 0; i < sample_vec[0].size(); i++) {
+            CHECK((sample_vec[0][i] == sample_vec[1][i]));
+        }
     }
 }

pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp

@@ -14,6 +14,7 @@
 #pragma once
 #include <chrono>
 #include <cstdint>
+#include <optional>
 
 #include <Kokkos_Core.hpp>
 #include <Kokkos_Random.hpp>
@@ -649,13 +650,16 @@ class Measurements final
      * Reference https://en.wikipedia.org/wiki/Inverse_transform_sampling
      *
      * @param num_samples Number of Samples
+     * @param seed Seed to generate the samples from
      *
      * @return std::vector<std::size_t> to the samples.
      * Each sample has a length equal to the number of qubits. Each sample can
      * be accessed using the stride sample_id*num_qubits, where sample_id is a
      * number between 0 and num_samples-1.
      */
-    auto generate_samples(std::size_t num_samples) -> std::vector<std::size_t> {
+    auto generate_samples(std::size_t num_samples,
+                          const std::optional<std::size_t> &seed = std::nullopt)
+        -> std::vector<std::size_t> {
         const std::size_t num_qubits = this->_statevector.getNumQubits();
         const std::size_t N = this->_statevector.getLength();
         Kokkos::View<std::size_t *> samples("num_samples",
@@ -674,10 +678,12 @@ class Measurements final
             });
 
         // Sampling using Random_XorShift64_Pool
-        Kokkos::Random_XorShift64_Pool<> rand_pool(
-            std::chrono::high_resolution_clock::now()
-                .time_since_epoch()
-                .count());
+        auto rand_pool = seed.has_value()
+                             ? Kokkos::Random_XorShift64_Pool<>(seed.value())
+                             : Kokkos::Random_XorShift64_Pool<>(
+                                   std::chrono::high_resolution_clock::now()
+                                       .time_since_epoch()
+                                       .count());
 
         Kokkos::parallel_for(
             Kokkos::RangePolicy<KokkosExecSpace>(0, num_samples),

pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp

@@ -23,6 +23,7 @@
 #include <algorithm>
 #include <complex>
 #include <cstdio>
+#include <optional>
 #include <random>
 #include <type_traits>
 #include <unordered_map>
@@ -573,30 +574,39 @@ class Measurements final
      * Reference: https://en.wikipedia.org/wiki/Alias_method
      *
      * @param num_samples The number of samples to generate.
+     * @param seed Seed to generate the samples from
      * @return 1-D vector of samples in binary, each sample is
      * separated by a stride equal to the number of qubits.
      */
-    std::vector<std::size_t> generate_samples(const std::size_t num_samples) {
+    std::vector<std::size_t>
+    generate_samples(const std::size_t num_samples,
+                     const std::optional<std::size_t> &seed = std::nullopt) {
         const std::size_t num_qubits = this->_statevector.getNumQubits();
         std::vector<std::size_t> wires(num_qubits);
         std::iota(wires.begin(), wires.end(), 0);
-        return generate_samples(wires, num_samples);
+        return generate_samples(wires, num_samples, seed);
     }
 
     /**
      * @brief Generate samples.
      *
      * @param wires Sample are generated for the specified wires.
      * @param num_samples The number of samples to generate.
+     * @param seed Seed to generate the samples from
      * @return 1-D vector of samples in binary, each sample is
      * separated by a stride equal to the number of qubits.
      */
     std::vector<std::size_t>
     generate_samples(const std::vector<std::size_t> &wires,
-                     const std::size_t num_samples) {
+                     const std::size_t num_samples,
+                     const std::optional<std::size_t> &seed = std::nullopt) {
         const std::size_t n_wires = wires.size();
         std::vector<std::size_t> samples(num_samples * n_wires);
-        this->setRandomSeed();
+        if (seed.has_value()) {
+            this->setSeed(seed.value());
+        } else {
+            this->setRandomSeed();
+        }
         DiscreteRandomVariable<PrecisionT> drv{this->rng, probs(wires)};
         // The Python layer expects a 2D array with dimensions (n_samples x
         // n_wires) and hence the linear index is `s * n_wires + (n_wires - 1 -