The lightning kokkos files in the Catalyst repo has changed since #770.…

… We make a small update to track these changes. The Catalyst PR adds seeding to qjit: PennyLaneAI/catalyst#936
PennyLaneAI · Jul 26, 2024 · 32e3e7e · 32e3e7e
1 parent d1b4007
commit 32e3e7e
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 1 deletion.
diff --git a/...lane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp b/...lane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp
@@ -106,6 +106,10 @@ auto LightningKokkosSimulator::GetDeviceShots() const -> std::size_t {
     return this->device_shots;
 }
 
+void LightningKokkosSimulator::SetDevicePRNG(std::mt19937 *gen) {
+    this->gen = gen;
+}
+
 /// LCOV_EXCL_START
 void LightningKokkosSimulator::PrintState() {
     using std::cout;
@@ -480,7 +484,7 @@ auto LightningKokkosSimulator::Measure(QubitIdType wire,
     SetDeviceShots(device_shots);
 
     // It represents the measured result, true for 1, false for 0
-    bool mres = Lightning::simulateDraw(probs, postselect);
+    bool mres = Lightning::simulateDraw(probs, postselect, this->gen);
     auto dev_wires = getDeviceWires(wires);
     this->device_sv->collapse(dev_wires[0], mres ? 1 : 0);
     return mres ? this->One() : this->Zero();

diff --git a/...lane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.hpp b/...lane_lightning/core/src/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.hpp
@@ -63,6 +63,8 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
 
     std::size_t device_shots;
 
+    std::mt19937 *gen = nullptr;
+
     std::unique_ptr<StateVectorT> device_sv = std::make_unique<StateVectorT>(0);
     LightningKokkosObsManager<double> obs_manager{};
 
@@ -116,6 +118,7 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
     void StartTapeRecording() override;
     void StopTapeRecording() override;
     void SetDeviceShots(size_t shots) override;
+    void SetDevicePRNG(std::mt19937 *) override;
     [[nodiscard]] auto GetDeviceShots() const -> size_t override;
     void PrintState() override;
     [[nodiscard]] auto Zero() const -> Result override;

diff --git a/...ning/core/src/simulators/lightning_kokkos/catalyst/tests/Test_LightningKokkosMeasures.cpp b/...ning/core/src/simulators/lightning_kokkos/catalyst/tests/Test_LightningKokkosMeasures.cpp
@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <random>
+
 #include "CacheManager.hpp"
 #include "LightningKokkosSimulator.hpp"
 #include "QuantumDevice.hpp"
@@ -1750,3 +1752,28 @@ TEST_CASE("Counts and PartialCounts tests with numWires=0-4 shots=100",
     CHECK(sum3 == shots);
     CHECK(sum4 == shots);
 }
+
+TEST_CASE("Measurement with a seeded device", "[Measures]") {
+    for (size_t _ = 0; _ < 5; _++) {
+        std::unique_ptr<LKSimulator> sim = std::make_unique<LKSimulator>();
+        std::unique_ptr<LKSimulator> sim1 = std::make_unique<LKSimulator>();
+
+        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);
+        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);
+    }
+}