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

pennylane_lightning/core/_version.py

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

pennylane_lightning/core/src/measurements/MeasurementsBase.hpp

@@ -77,13 +77,19 @@ 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;
         setSeed(rd());
     }
 
+    /**
+     * @brief Set the internal random generator to an already existing instance
+     *
+     * @param rng An already existing instance of a random number generator
+     */
+    void setRNG(std::mt19937 rng) { this->rng = rng; }
+
     /**
      * @brief Calculate the expectation value for a general Observable.
      *

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

@@ -20,6 +20,8 @@ using Pennylane::Util::isApproxEqual;
 } // namespace
 /// @endcond
 #include <algorithm>
+#include <optional>
+#include <random>
 #include <string>
 
 #ifdef _ENABLE_PLQUBIT
@@ -1251,7 +1253,9 @@ TEST_CASE("Var Shot- TensorProdObs", "[MeasurementsBase][Observables]") {
         testTensorProdObsVarShot<TestStateVectorBackends>();
     }
 }
-template <typename TypeList> void testSamples() {
+
+template <typename TypeList>
+void testSamples(const std::optional<std::mt19937> &rng = std::nullopt) {
     if constexpr (!std::is_same_v<TypeList, void>) {
         using StateVectorT = typename TypeList::Type;
         using PrecisionT = typename StateVectorT::PrecisionT;
@@ -1281,7 +1285,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 =
+            rng.has_value()
+                ? Measurer.generate_samples(num_samples, rng.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 +1314,7 @@ template <typename TypeList> void testSamples() {
             REQUIRE_THAT(probabilities,
                          Catch::Approx(expected_probabilities).margin(.05));
         }
-        testSamples<typename TypeList::Next>();
+        testSamples<typename TypeList::Next>(rng);
     }
 }
 
@@ -1317,6 +1324,13 @@ TEST_CASE("Samples", "[MeasurementsBase]") {
     }
 }
 
+TEST_CASE("Seeded samples", "[MeasurementsBase]") {
+    if constexpr (BACKEND_FOUND) {
+        std::mt19937 rng(37);
+        testSamples<TestStateVectorBackends>(rng);
+    }
+}
+
 template <typename TypeList> void testSamplesCountsObs() {
     if constexpr (!std::is_same_v<TypeList, void>) {
         using StateVectorT = typename TypeList::Type;
@@ -1729,4 +1743,4 @@ TEST_CASE("Measure Shot - SparseHObs ", "[MeasurementsBase][Observables]") {
     if constexpr (BACKEND_FOUND) {
         testSparseHObsMeasureShot<TestStateVectorBackends>();
     }
-}
+}

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,38 @@ class Measurements final
      * Reference: https://en.wikipedia.org/wiki/Alias_method
      *
      * @param num_samples The number of samples to generate.
+     * @param catalyst_rng The seeded random number generator from catalyst.
      * @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::mt19937> &catalyst_rng = 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, catalyst_rng);
     }
 
     /**
      * @brief Generate samples.
      *
      * @param wires Sample are generated for the specified wires.
      * @param num_samples The number of samples to generate.
+     * @param catalyst_rng The seeded random number generator from catalyst.
      * @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) {
+    std::vector<std::size_t> generate_samples(
+        const std::vector<std::size_t> &wires, const std::size_t num_samples,
+        const std::optional<std::mt19937> &catalyst_rng = std::nullopt) {
         const std::size_t n_wires = wires.size();
         std::vector<std::size_t> samples(num_samples * n_wires);
-        this->setRandomSeed();
+        if (catalyst_rng.has_value()) {
+            this->setRNG(catalyst_rng.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 -