Generalize seeding for measurements

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.40.0-dev14"
+__version__ = "0.40.0-dev15"

pennylane_lightning/core/src/measurements/MeasurementsBase.hpp

@@ -388,7 +388,8 @@ template <class StateVectorT, class Derived> class MeasurementsBase {
      */
     auto sample(const std::size_t &num_shots) -> std::vector<std::size_t> {
         Derived measure(_statevector);
-        return measure.generate_samples(num_shots);
+        return measure.generate_samples(num_shots,
+                                        static_cast<Derived *>(this));
     }
 
     /**
@@ -479,12 +480,14 @@ template <class StateVectorT, class Derived> class MeasurementsBase {
             StateVectorT sv(data_storage.data(), data_storage.size());
             obs.applyInPlaceShots(sv, eigenValues, obs_wires);
             Derived measure(sv);
-            samples = measure.generate_samples(num_shots);
+            samples = measure.generate_samples(num_shots,
+                                               static_cast<Derived *>(this));
         } else {
             StateVectorT sv(_statevector);
             obs.applyInPlaceShots(sv, eigenValues, obs_wires);
             Derived measure(sv);
-            samples = measure.generate_samples(num_shots);
+            samples = measure.generate_samples(num_shots,
+                                               static_cast<Derived *>(this));
         }
 
         if (!shot_range.empty()) {

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

@@ -1235,7 +1235,7 @@ TEST_CASE("Var Shot- TensorProdObs", "[MeasurementsBase][Observables]") {
 }
 
 template <typename TypeList>
-void testSamples(const std::optional<std::size_t> &seed = std::nullopt) {
+void testSamples(const std::optional<std::mt19937 *> &gen = std::nullopt) {
     if constexpr (!std::is_same_v<TypeList, void>) {
         using StateVectorT = typename TypeList::Type;
         using PrecisionT = typename StateVectorT::PrecisionT;
@@ -1257,6 +1257,9 @@ void testSamples(const std::optional<std::size_t> &seed = std::nullopt) {
         // This object attaches to the statevector allowing several
         // measurements.
         Measurements<StateVectorT> Measurer(statevector);
+        if (gen.has_value()) {
+            Measurer.set_PRNG(gen.value());
+        }
 
         std::vector<PrecisionT> expected_probabilities = {
             0.67078706, 0.03062806, 0.0870997,  0.00397696,
@@ -1265,10 +1268,7 @@ void testSamples(const std::optional<std::size_t> &seed = std::nullopt) {
         std::size_t num_qubits = 3;
         std::size_t N = std::pow(2, num_qubits);
         std::size_t num_samples = 100000;
-        auto &&samples =
-            seed.has_value()
-                ? Measurer.generate_samples(num_samples, seed.value())
-                : Measurer.generate_samples(num_samples);
+        auto &&samples = Measurer.generate_samples(num_samples);
 
         std::vector<std::size_t> counts(N, 0);
         std::vector<std::size_t> samples_decimal(num_samples, 0);
@@ -1294,7 +1294,7 @@ void testSamples(const std::optional<std::size_t> &seed = std::nullopt) {
             REQUIRE_THAT(probabilities,
                          Catch::Approx(expected_probabilities).margin(.05));
         }
-        testSamples<typename TypeList::Next>(seed);
+        testSamples<typename TypeList::Next>(gen);
     }
 }
 
@@ -1306,7 +1306,8 @@ TEST_CASE("Samples", "[MeasurementsBase]") {
 
 TEST_CASE("Seeded samples", "[MeasurementsBase]") {
     if constexpr (BACKEND_FOUND) {
-        testSamples<TestStateVectorBackends>(37);
+        std::mt19937 gen(37);
+        testSamples<TestStateVectorBackends>(&gen);
     }
 }
 

pennylane_lightning/core/src/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

@@ -244,6 +244,8 @@ auto LightningGPUSimulator::Expval(ObsIdType obsKey) -> double {
     Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     return device_shots ? m.expval(*obs, device_shots, {}) : m.expval(*obs);
 }
 
@@ -261,6 +263,8 @@ auto LightningGPUSimulator::Var(ObsIdType obsKey) -> double {
     Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     return device_shots ? m.var(*obs, device_shots) : m.var(*obs);
 }
 
@@ -282,6 +286,9 @@ void LightningGPUSimulator::State(DataView<std::complex<double>, 1> &state) {
 void LightningGPUSimulator::Probs(DataView<double, 1> &probs) {
     Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
+
+    m.set_PRNG(this->gen);
+
     auto &&dv_probs = device_shots ? m.probs(device_shots) : m.probs();
 
     RT_FAIL_IF(probs.size() != dv_probs.size(),
@@ -301,6 +308,9 @@ void LightningGPUSimulator::PartialProbs(
     auto dev_wires = getDeviceWires(wires);
     Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
+
+    m.set_PRNG(this->gen);
+
     auto &&dv_probs =
         device_shots ? m.probs(dev_wires, device_shots) : m.probs(dev_wires);
 
@@ -315,9 +325,8 @@ std::vector<size_t> LightningGPUSimulator::GenerateSamples(size_t shots) {
     Pennylane::LightningGPU::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
-    if (this->gen) {
-        return m.generate_samples(shots, (*(this->gen))());
-    }
+    m.set_PRNG(this->gen);
+
     return m.generate_samples(shots);
 }
 

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

@@ -73,6 +73,8 @@ class Measurements final
 
     GateCache<PrecisionT> gate_cache_;
 
+    std::mt19937 *gen = nullptr;
+
   public:
     explicit Measurements(StateVectorT &statevector)
         : BaseType{statevector},
@@ -215,9 +217,7 @@ 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,
-                          const std::optional<std::size_t> &seed = std::nullopt)
-        -> std::vector<std::size_t> {
+    auto generate_samples(std::size_t num_samples) -> std::vector<std::size_t> {
         std::vector<double> rand_nums(num_samples);
         custatevecSamplerDescriptor_t sampler;
 
@@ -237,8 +237,9 @@ class Measurements final
             data_type = CUDA_C_32F;
         }
 
-        if (seed.has_value()) {
-            this->setSeed(seed.value());
+        if (gen != nullptr) {
+            std::size_t seed = (*(this->gen))();
+            this->setSeed(seed);
         } else {
             this->setRandomSeed();
         }
@@ -713,6 +714,8 @@ class Measurements final
         return BaseType::var(obs, num_shots);
     }
 
+    void set_PRNG(std::mt19937 *gen) { this->gen = gen; }
+
   private:
     /**
      * @brief Utility method for expectation value calculations.

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

@@ -268,6 +268,8 @@ auto LightningKokkosSimulator::Expval(ObsIdType obsKey) -> double {
     Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     return device_shots ? m.expval(*obs, device_shots, {}) : m.expval(*obs);
 }
 
@@ -285,6 +287,8 @@ auto LightningKokkosSimulator::Var(ObsIdType obsKey) -> double {
     Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     return device_shots ? m.var(*obs, device_shots) : m.var(*obs);
 }
 
@@ -314,6 +318,9 @@ void LightningKokkosSimulator::State(DataView<std::complex<double>, 1> &state) {
 void LightningKokkosSimulator::Probs(DataView<double, 1> &probs) {
     Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
+
+    m.set_PRNG(this->gen);
+
     auto &&dv_probs = device_shots ? m.probs(device_shots) : m.probs();
 
     RT_FAIL_IF(probs.size() != dv_probs.size(),
@@ -333,6 +340,9 @@ void LightningKokkosSimulator::PartialProbs(
     auto dev_wires = getDeviceWires(wires);
     Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
+
+    m.set_PRNG(this->gen);
+
     auto &&dv_probs =
         device_shots ? m.probs(dev_wires, device_shots) : m.probs(dev_wires);
 
@@ -347,11 +357,10 @@ std::vector<size_t> LightningKokkosSimulator::GenerateSamples(size_t shots) {
     Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     // PL-Lightning-Kokkos generates samples using the alias method.
     // Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
-    if (this->gen) {
-        return m.generate_samples(shots, (*(this->gen))());
-    }
     return m.generate_samples(shots);
 }
 

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

@@ -15,6 +15,7 @@
 #include <chrono>
 #include <cstdint>
 #include <optional>
+#include <random>
 
 #include <Kokkos_Core.hpp>
 #include <Kokkos_Random.hpp>
@@ -72,6 +73,8 @@ class Measurements final
     using ScratchViewComplex = typename StateVectorT::ScratchViewComplex;
     using TeamPolicy = typename StateVectorT::TeamPolicy;
 
+    std::mt19937 *gen = nullptr;
+
   public:
     explicit Measurements(const StateVectorT &statevector)
         : BaseType{statevector} {
@@ -650,16 +653,13 @@ 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,
-                          const std::optional<std::size_t> &seed = std::nullopt)
-        -> std::vector<std::size_t> {
+    auto generate_samples(std::size_t num_samples) -> 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",
@@ -678,12 +678,13 @@ class Measurements final
             });
 
         // Sampling using Random_XorShift64_Pool
-        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());
+        auto rand_pool =
+            this->gen != nullptr
+                ? Kokkos::Random_XorShift64_Pool<>((*(this->gen))())
+                : Kokkos::Random_XorShift64_Pool<>(
+                      std::chrono::high_resolution_clock::now()
+                          .time_since_epoch()
+                          .count());
 
         Kokkos::parallel_for(
             Kokkos::RangePolicy<KokkosExecSpace>(0, num_samples),
@@ -693,6 +694,8 @@ class Measurements final
         return view2vector(samples);
     }
 
+    void set_PRNG(std::mt19937 *gen) { this->gen = gen; }
+
   private:
     std::unordered_map<std::string, ExpValFunc> expval_funcs_;
 

pennylane_lightning/core/src/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

@@ -251,6 +251,8 @@ auto LightningSimulator::Expval(ObsIdType obsKey) -> double {
     Pennylane::LightningQubit::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     return (device_shots != 0U) ? m.expval(*obs, device_shots, {})
                                 : m.expval(*obs);
 }
@@ -268,6 +270,8 @@ auto LightningSimulator::Var(ObsIdType obsKey) -> double {
     Pennylane::LightningQubit::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     return (device_shots != 0U) ? m.var(*obs, device_shots) : m.var(*obs);
 }
 
@@ -282,6 +286,9 @@ void LightningSimulator::State(DataView<std::complex<double>, 1> &state) {
 void LightningSimulator::Probs(DataView<double, 1> &probs) {
     Pennylane::LightningQubit::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
+
+    m.set_PRNG(this->gen);
+
     auto &&dv_probs = (device_shots != 0U) ? m.probs(device_shots) : m.probs();
 
     RT_FAIL_IF(probs.size() != dv_probs.size(),
@@ -301,6 +308,9 @@ void LightningSimulator::PartialProbs(DataView<double, 1> &probs,
     auto dev_wires = getDeviceWires(wires);
     Pennylane::LightningQubit::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
+
+    m.set_PRNG(this->gen);
+
     auto &&dv_probs = (device_shots != 0U) ? m.probs(dev_wires, device_shots)
                                            : m.probs(dev_wires);
 
@@ -316,6 +326,8 @@ LightningSimulator::GenerateSamplesMetropolis(size_t shots) {
     Pennylane::LightningQubit::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     // PL-Lightning generates samples using the alias method.
     // Reference: https://en.wikipedia.org/wiki/Alias_method
     // Given the number of samples, returns 1-D vector of samples
@@ -335,16 +347,15 @@ std::vector<size_t> LightningSimulator::GenerateSamples(size_t shots) {
     Pennylane::LightningQubit::Measures::Measurements<StateVectorT> m{
         *(this->device_sv)};
 
+    m.set_PRNG(this->gen);
+
     // PL-Lightning generates samples using the alias method.
     // Reference: https://en.wikipedia.org/wiki/Alias_method
     // Given the number of samples, returns 1-D vector of samples
     // in binary, each sample is separated by a stride equal to
     // the number of qubits.
     //
     // Return Value Optimization (RVO)
-    if (this->gen != nullptr) {
-        return m.generate_samples(shots, (*(this->gen))());
-    }
     return m.generate_samples(shots);
 }