Seed samples for lightning.qubit/kokkos (#1164)

**Context:**
There's still quite a few frontend tests stochastically failing because
the `seed` option in `qjit` only controls the measurements, but not the
samples. We add seeding to the samples.

**Description of the Change:**
When `qjit(seed=...)` receives a (unsigned 32 bit int) seed value from
the user, the seed gets propagated through mlir and [eventually becomes
a field of the `Catalyst::Runtime::Simulator::LightningSimulator` class,
alongside the seeded `std::mt19937` rng instance
](https://github.com/PennyLaneAI/catalyst/blob/a580bada575793b780d5366aa77dff6157cd4f93/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.hpp#L54).
This was done in #936.

In #936 , [the device's rng instance is used during measurements
](https://github.com/PennyLaneAI/catalyst/blob/a580bada575793b780d5366aa77dff6157cd4f93/runtime/lib/backend/lightning/lightning_dynamic/LightningSimulator.cpp#L451),
but not during samples. This is because samples are performed from the
`Pennylane::LightningQubit::Measures::Measurements` class through the
`generate_samples` methods, which is controlled by the lightning repo.

To seed samples, we use the device rng instance to generate a
deterministic seed to pass it onto the state vector's `generate_samples`
methods. This is the only change in catalyst.


In lightning, the `generate_samples` method now can take in a seeding
number. The catalyst devices pass in a seed into the lightning
`generate_samples`; this seed is created deterministically from the
aforementioned already seeded catalyst context rng instance. This makes
the generated samples deterministc. The above is published on the
lightning repo as the branch "seed_sample_lightning":
PennyLaneAI/pennylane-lightning#927

PennyLaneAI/pennylane-lightning@6f3e0d5

**Benefits:**
Fewer (hopefully no) stochatically failing frontend tests.


**Related GitHub Issues:** #999 
[sc-72878]

.dep-versions

@@ -15,4 +15,4 @@ pennylane=0.39.0.dev26
 # 'runtime/Makefile' and at all GitHub workflows, using the exact
 # commit hash corresponding to the merged PR that implements the
 # desired feature.
-lightning=0.38.0
+lightning=0.39.0-dev38

.github/workflows/build-wheel-linux-x86_64.yaml

@@ -342,7 +342,7 @@ jobs:
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
               -DENABLE_LAPACK=OFF \
-              -DLIGHTNING_GIT_TAG=latest_release \
+              -DLIGHTNING_GIT_TAG=6f3e0d5d371ff9823a3177dd2c66052668883d42 \
               -DENABLE_WARNINGS=OFF \
               -DENABLE_OPENQASM=ON \
               -DENABLE_OPENMP=OFF \

.github/workflows/build-wheel-macos-arm64.yaml

@@ -301,7 +301,7 @@ jobs:
               -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$GITHUB_WORKSPACE/runtime-build/lib \
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
-              -DLIGHTNING_GIT_TAG=latest_release \
+              -DLIGHTNING_GIT_TAG=6f3e0d5d371ff9823a3177dd2c66052668883d42\
               -DENABLE_LAPACK=OFF \
               -DENABLE_WARNINGS=OFF \
               -DENABLE_OPENQASM=ON \

.github/workflows/build-wheel-macos-x86_64.yaml

@@ -292,7 +292,7 @@ jobs:
               -DCMAKE_LIBRARY_OUTPUT_DIRECTORY=$GITHUB_WORKSPACE/runtime-build/lib \
               -DPYTHON_EXECUTABLE=$(which python${{ matrix.python_version }}) \
               -Dpybind11_DIR=$(python${{ matrix.python_version }} -c "import pybind11; print(pybind11.get_cmake_dir())") \
-              -DLIGHTNING_GIT_TAG=latest_release \
+              -DLIGHTNING_GIT_TAG=6f3e0d5d371ff9823a3177dd2c66052668883d42 \
               -DENABLE_LAPACK=OFF \
               -DENABLE_WARNINGS=OFF \
               -DENABLE_OPENQASM=ON \

.github/workflows/scripts/linux_arm64/rh8/build_catalyst.sh

@@ -45,7 +45,7 @@ cmake -S runtime -B runtime-build -G Ninja \
     -DPYTHON_INCLUDE_DIR=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/include/python${PYTHON_VERSION} \
     -DPYTHON_LIBRARY=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/lib \
     -Dpybind11_DIR=/opt/_internal/cpython-${PYTHON_VERSION}.${PYTHON_SUBVERSION}/lib/python${PYTHON_VERSION}/site-packages/pybind11/share/cmake/pybind11 \
-    -DLIGHTNING_GIT_TAG=latest_release \
+    -DLIGHTNING_GIT_TAG=6f3e0d5d371ff9823a3177dd2c66052668883d42 \
     -DENABLE_LAPACK=OFF \
     -DENABLE_WARNINGS=OFF \
     -DENABLE_OPENQASM=ON \

doc/releases/changelog-dev.md

@@ -202,6 +202,9 @@
 * Cached primitive lowerings is used instead of a custom cache structure.
   [(#1159)](https://github.com/PennyLaneAI/catalyst/pull/1159)
 
+* Samples on lightning.qubit/kokkos can now be seeded with `qjit(seed=...)`.
+  [(#1164)](https://github.com/PennyLaneAI/catalyst/pull/1164)
+
 <h3>Breaking changes</h3>
 
 * Remove `static_size` field from `AbstractQreg` class.

doc/requirements.txt

@@ -30,6 +30,6 @@ lxml_html_clean
 
 # Pre-install PL development wheels
 --extra-index-url https://test.pypi.org/simple/
-pennylane-lightning-kokkos==0.38.0
-pennylane-lightning==0.38.0
+pennylane-lightning-kokkos==0.39.0-dev38
+pennylane-lightning==0.39.0-dev38
 pennylane==0.39.0.dev26

frontend/test/pytest/test_measurement_transforms.py

@@ -167,7 +167,7 @@ def basic_circuit(theta: float):
         assert "counts" in mlir
 
         theta = 1.9
-        expval_res, var_res, counts_res, probs_res = qml.qjit(transformed_circuit)(theta)
+        expval_res, var_res, counts_res, probs_res = qml.qjit(transformed_circuit, seed=37)(theta)
 
         expval_expected = np.sin(theta) * np.sin(theta / 2)
         var_expected = 1 - np.sin(2 * theta) ** 2
@@ -221,7 +221,7 @@ def basic_circuit(theta: float):
 
         theta = 1.9
 
-        expval_res, var_res, sample_res, probs_res = qml.qjit(transformed_circuit)(theta)
+        expval_res, var_res, sample_res, probs_res = qml.qjit(transformed_circuit, seed=37)(theta)
 
         expval_expected = np.sin(theta) * np.sin(theta / 2)
         var_expected = 1 - np.sin(2 * theta) ** 2
@@ -389,7 +389,7 @@ def circuit(theta: float):
 
         theta = 2.5
         counts_expected = circuit(theta)
-        res = qml.qjit(measurements_from_counts(circuit, dev.wires))(theta)
+        res = qml.qjit(measurements_from_counts(circuit, dev.wires), seed=37)(theta)
 
         # counts comparison by converting catalyst format to PL style eigvals dict
         basis_states, counts = res
@@ -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 not in ["lightning.qubit", "lightning.kokkos"]:
+        pytest.skip("Sample seeding is only supported on lightning.qubit and lightning.kokkos")
+
+    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?=6f3e0d5d371ff9823a3177dd2c66052668883d42
 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);
 }
 

runtime/tests/Test_SVDynamicCPU_Allocation.cpp

@@ -24,7 +24,7 @@
 #include "LinearAlgebra.hpp"
 #include "StateVectorLQubitDynamic.hpp"
 #include "Util.hpp"
-#include "cpu_kernels/GateImplementationsPI.hpp"
+#include "cpu_kernels/GateImplementationsLM.hpp"
 #include <StateVectorLQubit.hpp>
 
 #include "TestHelpers.hpp"

runtime/tests/Test_SVDynamicCPU_Core.cpp

@@ -24,7 +24,7 @@
 #include "LinearAlgebra.hpp"
 #include "StateVectorLQubitDynamic.hpp"
 #include "Util.hpp"
-#include "cpu_kernels/GateImplementationsPI.hpp"
+#include "cpu_kernels/GateImplementationsLM.hpp"
 #include <StateVectorLQubit.hpp>
 
 #include "TestHelpers.hpp"
@@ -129,7 +129,7 @@ TEMPLATE_TEST_CASE("StateVectorLQubitDynamic::applyMatrix with a pointer",
 
             const auto m = Pennylane::Util::randomUnitary<PrecisionT>(re, num_wires);
             sv1.applyMatrix(m, wires);
-            Gates::GateImplementationsPI::applyMultiQubitOp<PrecisionT>(sv2.getData(), num_qubits,
+            Gates::GateImplementationsLM::applyMultiQubitOp<PrecisionT>(sv2.getData(), num_qubits,
                                                                         m.data(), wires, false);
             CHECK(sv1.getDataVector() == approx(sv2.getDataVector()).margin(PrecisionT{1e-5}));
         }