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

.github/CHANGELOG.md

@@ -37,6 +37,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)
 
@@ -105,7 +108,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-dev35"
+__version__ = "0.39.0-dev36"

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,23 @@ 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(const 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,
+                                      const 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");
@@ -369,7 +379,7 @@ void LightningKokkosSimulator::Sample(DataView<double, 2> &samples,
 }
 void LightningKokkosSimulator::PartialSample(
     DataView<double, 2> &samples, const std::vector<QubitIdType> &wires,
-    std::size_t shots) {
+    const std::size_t shots) {
     const std::size_t numWires = wires.size();
     const std::size_t numQubits = this->GetNumQubits();
 
@@ -381,13 +391,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
@@ -404,20 +408,14 @@ void LightningKokkosSimulator::PartialSample(
 
 void LightningKokkosSimulator::Counts(DataView<double, 1> &eigvals,
                                       DataView<int64_t, 1> &counts,
-                                      std::size_t shots) {
+                                      const std::size_t shots) {
     const std::size_t numQubits = this->GetNumQubits();
     const std::size_t numElements = 1U << numQubits;
 
     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
@@ -442,7 +440,7 @@ void LightningKokkosSimulator::Counts(DataView<double, 1> &eigvals,
 
 void LightningKokkosSimulator::PartialCounts(
     DataView<double, 1> &eigvals, DataView<int64_t, 1> &counts,
-    const std::vector<QubitIdType> &wires, std::size_t shots) {
+    const std::vector<QubitIdType> &wires, const std::size_t shots) {
     const std::size_t numWires = wires.size();
     const std::size_t numQubits = this->GetNumQubits();
     const std::size_t numElements = 1U << numWires;
@@ -455,13 +453,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(const size_t shots) -> std::vector<size_t>;
+
   public:
     explicit LightningKokkosSimulator(const std::string &kwargs = "{}") {
         auto &&args = Catalyst::Runtime::parse_kwargs(kwargs);
@@ -155,16 +157,16 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
     void Probs(DataView<double, 1> &probs) override;
     void PartialProbs(DataView<double, 1> &probs,
                       const std::vector<QubitIdType> &wires) override;
-    void Sample(DataView<double, 2> &samples, std::size_t shots) override;
+    void Sample(DataView<double, 2> &samples, const std::size_t shots) override;
     void PartialSample(DataView<double, 2> &samples,
                        const std::vector<QubitIdType> &wires,
-                       std::size_t shots) override;
+                       const std::size_t shots) override;
     void Counts(DataView<double, 1> &eigvals, DataView<int64_t, 1> &counts,
-                std::size_t shots) override;
+                const std::size_t shots) override;
     void PartialCounts(DataView<double, 1> &eigvals,
                        DataView<int64_t, 1> &counts,
                        const std::vector<QubitIdType> &wires,
-                       std::size_t shots) override;
+                       const std::size_t shots) override;
     auto Measure(QubitIdType wire,
                  std::optional<int32_t> postselect = std::nullopt)
         -> Result override;

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

@@ -1777,3 +1777,46 @@ TEST_CASE("Measurement with a seeded device", "[Measures]") {
         CHECK(*m == *m1);
     }
 }
+
+TEST_CASE("Sample 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::size_t shots = 100;
+
+        std::vector<double> samples(shots * 4);
+        MemRefT<double, 2> buffer{
+            samples.data(), samples.data(), 0, {shots, 1}, {1, 1}};
+        DataView<double, 2> view(buffer.data_aligned, buffer.offset,
+                                 buffer.sizes, buffer.strides);
+
+        std::vector<double> samples1(shots * 4);
+        MemRefT<double, 2> buffer1{
+            samples1.data(), samples1.data(), 0, {shots, 1}, {1, 1}};
+        DataView<double, 2> view1(buffer1.data_aligned, buffer1.offset,
+                                  buffer1.sizes, buffer1.strides);
+
+        std::mt19937 gen(37);
+        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);
+
+        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);
+        sim1->NamedOperation("RX", {0.5}, {Qs1[0]}, false);
+        sim1->Sample(view1, shots);
+
+        for (std::size_t i = 0; i < shots * 4; i++) {
+            CHECK((samples[i] == samples1[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 -