Seed samples for lightning.qubit/kokkos

frontend/test/pytest/test_measurement_transforms.py

@@ -441,7 +441,7 @@ def circuit(theta: float):
             return measurement()
 
         theta = 2.5
-        res = qml.qjit(measurements_from_samples(circuit, dev.wires))(theta)
+        res = qml.qjit(measurements_from_samples(circuit, dev.wires), seed=37)(theta)
 
         if len(measurement().wires) == 1:
             samples_expected = qml.qjit(circuit)(theta)

frontend/test/pytest/test_seeded_qjit.py

@@ -99,5 +99,53 @@ def cfun0():
         assert np.allclose(results0, results2)
 
 
+@pytest.mark.parametrize(
+    "seed",
+    [
+        42,
+        37,
+        1337,
+        2**32 - 1,
+        0,
+    ],
+)
+@pytest.mark.parametrize("shots", [10])
+def test_seeded_sample(seed, shots, backend):
+    """Test that different calls to qjits with the same seed produce the same sample results"""
+
+    if backend != "lightning.qubit":
+        pytest.skip("Sample seeding is only supported on lightning.qubit")
+
+    dev = qml.device(backend, wires=2, shots=shots)
+
+    @qjit(seed=seed)
+    def workflow():
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=[0])
+            qml.RX(12.34, wires=[1])
+            return qml.sample()
+
+        return circuit(), circuit(), circuit(), circuit()
+
+    @qjit(seed=seed)
+    def workflow1():
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=[0])
+            qml.RX(12.34, wires=[1])
+            return qml.sample()
+
+        return circuit(), circuit(), circuit(), circuit()
+
+    # Calls to qjits with the same seed should return the same samples
+    for _ in range(5):
+        results0 = workflow()
+        results1 = workflow()
+        results2 = workflow1()
+        assert np.allclose(results0, results1)
+        assert np.allclose(results0, results2)
+
+
 if __name__ == "__main__":
     pytest.main(["-x", __file__])

runtime/Makefile

@@ -14,7 +14,7 @@ ENABLE_LIGHTNING?=ON
 ENABLE_LIGHTNING_KOKKOS?=ON
 ENABLE_OPENQASM?=ON
 ENABLE_ASAN?=OFF
-LIGHTNING_GIT_TAG_VALUE?=latest_release
+LIGHTNING_GIT_TAG_VALUE?=seed_sample_lightning
 ENABLE_LAPACK?=OFF
 
 BUILD_TARGETS := rt_capi rtd_dummy

runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.cpp

@@ -314,6 +314,9 @@ std::vector<size_t> LightningSimulator::GenerateSamples(size_t shots)
     // the number of qubits.
     //
     // Return Value Optimization (RVO)
+    if (this->gen) {
+        return m.generate_samples(shots, *(this->gen));
+    }
     return m.generate_samples(shots);
 }