Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 39 commits into from
Oct 3, 2024
Merged
Show file tree
Hide file tree
Changes from 31 commits
Commits
Show all changes
39 commits
Select commit Hold shift + click to select a range
290a497
try getting catalyst's rng for generating samples
paul0403 Sep 30, 2024
133bc29
set default to nullptr
paul0403 Sep 30, 2024
d0f0777
Auto update version from '0.39.0-dev33' to '0.39.0-dev34'
ringo-but-quantum Sep 30, 2024
ea8290b
make the other `generate_samples` take in the rng as well
paul0403 Sep 30, 2024
ce4c532
undo kokkos changes for now
paul0403 Sep 30, 2024
9e5eb33
remove unnecessary `this->`
paul0403 Sep 30, 2024
6aa6d47
passing pointers as const reference
paul0403 Sep 30, 2024
cfb16b4
format
paul0403 Sep 30, 2024
3047599
convert to std::optional to avoid (pointers + default value of nullptr);
paul0403 Sep 30, 2024
0b943f1
give std::nullopt default to std::optional
paul0403 Sep 30, 2024
f496df8
test
paul0403 Oct 1, 2024
e116f5c
Auto update version from '0.39.0-dev34' to '0.39.0-dev35'
ringo-but-quantum Oct 1, 2024
3c52db6
Merge remote-tracking branch 'origin/master' into seed_sample_lightning
paul0403 Oct 1, 2024
5216cca
apply clang tidy comments
paul0403 Oct 1, 2024
234ced7
format
paul0403 Oct 1, 2024
c4c5341
pass the seeding integer, instead of the rng instance; add kokkos
paul0403 Oct 1, 2024
2a2420e
codefactor
paul0403 Oct 1, 2024
671fed4
debug push
paul0403 Oct 1, 2024
1f861df
debug
paul0403 Oct 1, 2024
3624cb0
remove debug prints
paul0403 Oct 2, 2024
5896119
small comments
paul0403 Oct 2, 2024
b2e3759
changelog
paul0403 Oct 2, 2024
9e98e33
changelog author
paul0403 Oct 2, 2024
e57ada4
`auto` for kokkos random pool type when type is obvious;
paul0403 Oct 2, 2024
f082ea1
catalyst kokkos seed sample test
paul0403 Oct 2, 2024
5edfdfa
format
paul0403 Oct 2, 2024
3fca3d4
Auto update version from '0.39.0-dev35' to '0.39.0-dev36'
ringo-but-quantum Oct 2, 2024
c2f8e1a
Merge remote-tracking branch 'origin/master' into seed_sample_lightning
paul0403 Oct 2, 2024
5101ff4
change helper to private
paul0403 Oct 2, 2024
a127f33
fix lightning.gpu test failure
paul0403 Oct 2, 2024
c70de5d
optional header
paul0403 Oct 2, 2024
82a3ed0
Merge remote-tracking branch 'origin/master' into seed_sample_lightning
paul0403 Oct 3, 2024
3be52d6
factor out common furnitures for all loop iterations for kokkos catal…
paul0403 Oct 3, 2024
92b64f7
remove const for size_t
paul0403 Oct 3, 2024
edfa8fd
rewrite the measurement seeding kokkos test to reuse common furniture…
paul0403 Oct 3, 2024
5306829
Auto update version from '0.39.0-dev36' to '0.39.0-dev37'
ringo-but-quantum Oct 3, 2024
ddc6732
small bug
paul0403 Oct 3, 2024
3362dd0
Merge remote-tracking branch 'origin/master' into seed_sample_lightning
paul0403 Oct 3, 2024
c4d67f4
Auto update version from '0.39.0-dev37' to '0.39.0-dev38'
ringo-but-quantum Oct 3, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion .github/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,9 @@

### Improvements

* The `generate_samples` methods of lightning.{qubit/kokkos} can now take in a seed number to make the generated samples deterministic. This can be useful when, among other things, fixing flaky tests in CI.
paul0403 marked this conversation as resolved.
Show resolved Hide resolved
[(#927)](https://github.com/PennyLaneAI/pennylane-lightning/pull/927)

* Always decompose `qml.QFT` in Lightning.
[(#924)](https://github.com/PennyLaneAI/pennylane-lightning/pull/924)

Expand Down Expand Up @@ -105,7 +108,7 @@

This release contains contributions from (in alphabetical order):

Ali Asadi, Amintor Dusko, Luis Alfredo Nuñez Meneses, Vincent Michaud-Rioux, Lee J. O'Riordan, Mudit Pandey, Shuli Shu
Ali Asadi, Amintor Dusko, Luis Alfredo Nuñez Meneses, Vincent Michaud-Rioux, Lee J. O'Riordan, Mudit Pandey, Shuli Shu, Haochen Paul Wang

---

Expand Down
2 changes: 1 addition & 1 deletion pennylane_lightning/core/_version.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,4 +16,4 @@
Version number (major.minor.patch[-label])
"""

__version__ = "0.39.0-dev35"
__version__ = "0.39.0-dev36"
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,6 @@ 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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ using Pennylane::Util::isApproxEqual;
} // namespace
/// @endcond
#include <algorithm>
#include <optional>
#include <string>

#ifdef _ENABLE_PLQUBIT
Expand Down Expand Up @@ -1251,7 +1252,9 @@ TEST_CASE("Var Shot- TensorProdObs", "[MeasurementsBase][Observables]") {
testTensorProdObsVarShot<TestStateVectorBackends>();
}
}
template <typename TypeList> void testSamples() {

template <typename TypeList>
void testSamples(const std::optional<std::size_t> &seed = std::nullopt) {
if constexpr (!std::is_same_v<TypeList, void>) {
using StateVectorT = typename TypeList::Type;
using PrecisionT = typename StateVectorT::PrecisionT;
Expand Down Expand Up @@ -1281,7 +1284,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 =
seed.has_value()
? Measurer.generate_samples(num_samples, seed.value())
: Measurer.generate_samples(num_samples);

std::vector<std::size_t> counts(N, 0);
std::vector<std::size_t> samples_decimal(num_samples, 0);
Expand All @@ -1307,7 +1313,7 @@ template <typename TypeList> void testSamples() {
REQUIRE_THAT(probabilities,
Catch::Approx(expected_probabilities).margin(.05));
}
testSamples<typename TypeList::Next>();
testSamples<typename TypeList::Next>(seed);
}
}

Expand All @@ -1317,6 +1323,12 @@ TEST_CASE("Samples", "[MeasurementsBase]") {
}
}

TEST_CASE("Seeded samples", "[MeasurementsBase]") {
if constexpr (BACKEND_FOUND) {
testSamples<TestStateVectorBackends>(37);
}
}

template <typename TypeList> void testSamplesCountsObs() {
if constexpr (!std::is_same_v<TypeList, void>) {
using StateVectorT = typename TypeList::Type;
Expand Down Expand Up @@ -1729,4 +1741,4 @@ TEST_CASE("Measure Shot - SparseHObs ", "[MeasurementsBase][Observables]") {
if constexpr (BACKEND_FOUND) {
testSparseHObsMeasureShot<TestStateVectorBackends>();
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
#include <cuda.h>
#include <cusparse.h>
#include <custatevec.h> // custatevecApplyMatrix
#include <optional>
#include <random>
#include <type_traits>
#include <unordered_map>
Expand Down Expand Up @@ -218,7 +219,9 @@ 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) -> std::vector<std::size_t> {
auto generate_samples(std::size_t num_samples,
const std::optional<std::size_t> &seed = std::nullopt)
-> std::vector<std::size_t> {
std::vector<double> rand_nums(num_samples);
custatevecSamplerDescriptor_t sampler;

Expand All @@ -238,7 +241,11 @@ class Measurements final
data_type = CUDA_C_32F;
}

this->setRandomSeed();
if (seed.has_value()) {
this->setSeed(seed.value());
} else {
this->setRandomSeed();
}
std::uniform_real_distribution<PrecisionT> dis(0.0, 1.0);
for (std::size_t n = 0; n < num_samples; n++) {
rand_nums[n] = dis(this->rng);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -342,13 +342,23 @@ void LightningKokkosSimulator::PartialProbs(
std::move(dv_probs.begin(), dv_probs.end(), probs.begin());
}

void LightningKokkosSimulator::Sample(DataView<double, 2> &samples,
std::size_t shots) {
std::vector<size_t>
LightningKokkosSimulator::GenerateSamples(const size_t shots) {
// generate_samples is a member function of the Measures class.
Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
*(this->device_sv)};

// PL-Lightning-Kokkos generates samples using the alias method.
// Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
auto li_samples = m.generate_samples(shots);
if (this->gen) {
return m.generate_samples(shots, (*(this->gen))());
paul0403 marked this conversation as resolved.
Show resolved Hide resolved
}
return m.generate_samples(shots);
paul0403 marked this conversation as resolved.
Show resolved Hide resolved
}

void LightningKokkosSimulator::Sample(DataView<double, 2> &samples,
const std::size_t shots) {
auto li_samples = this->GenerateSamples(shots);

RT_FAIL_IF(samples.size() != li_samples.size(),
"Invalid size for the pre-allocated samples");
Expand All @@ -369,7 +379,7 @@ void LightningKokkosSimulator::Sample(DataView<double, 2> &samples,
}
void LightningKokkosSimulator::PartialSample(
DataView<double, 2> &samples, const std::vector<QubitIdType> &wires,
std::size_t shots) {
const std::size_t shots) {
const std::size_t numWires = wires.size();
const std::size_t numQubits = this->GetNumQubits();

Expand All @@ -381,13 +391,7 @@ void LightningKokkosSimulator::PartialSample(
// get device wires
auto &&dev_wires = getDeviceWires(wires);

// generate_samples is a member function of the MeasuresKokkos class.
Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
*(this->device_sv)};

// PL-Lightning-Kokkos generates samples using the alias method.
// Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
auto li_samples = m.generate_samples(shots);
auto li_samples = this->GenerateSamples(shots);

// The lightning samples are layed out as a single vector of size
// shots*qubits, where each element represents a single bit. The
Expand All @@ -404,20 +408,14 @@ void LightningKokkosSimulator::PartialSample(

void LightningKokkosSimulator::Counts(DataView<double, 1> &eigvals,
DataView<int64_t, 1> &counts,
std::size_t shots) {
const std::size_t shots) {
mlxd marked this conversation as resolved.
Show resolved Hide resolved
const std::size_t numQubits = this->GetNumQubits();
const std::size_t numElements = 1U << numQubits;

RT_FAIL_IF(eigvals.size() != numElements || counts.size() != numElements,
"Invalid size for the pre-allocated counts");

// generate_samples is a member function of the MeasuresKokkos class.
Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
*(this->device_sv)};

// PL-Lightning-Kokkos generates samples using the alias method.
// Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
auto li_samples = m.generate_samples(shots);
auto li_samples = this->GenerateSamples(shots);

// Fill the eigenvalues with the integer representation of the corresponding
// computational basis bitstring. In the future, eigenvalues can also be
Expand All @@ -442,7 +440,7 @@ void LightningKokkosSimulator::Counts(DataView<double, 1> &eigvals,

void LightningKokkosSimulator::PartialCounts(
DataView<double, 1> &eigvals, DataView<int64_t, 1> &counts,
const std::vector<QubitIdType> &wires, std::size_t shots) {
const std::vector<QubitIdType> &wires, const std::size_t shots) {
const std::size_t numWires = wires.size();
const std::size_t numQubits = this->GetNumQubits();
const std::size_t numElements = 1U << numWires;
Expand All @@ -455,13 +453,7 @@ void LightningKokkosSimulator::PartialCounts(
// get device wires
auto &&dev_wires = getDeviceWires(wires);

// generate_samples is a member function of the MeasuresKokkos class.
Pennylane::LightningKokkos::Measures::Measurements<StateVectorT> m{
*(this->device_sv)};

// PL-Lightning-Kokkos generates samples using the alias method.
// Reference: https://en.wikipedia.org/wiki/Inverse_transform_sampling
auto li_samples = m.generate_samples(shots);
auto li_samples = this->GenerateSamples(shots);

// Fill the eigenvalues with the integer representation of the corresponding
// computational basis bitstring. In the future, eigenvalues can also be
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,8 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
return res;
}

auto GenerateSamples(const size_t shots) -> std::vector<size_t>;

public:
explicit LightningKokkosSimulator(const std::string &kwargs = "{}") {
auto &&args = Catalyst::Runtime::parse_kwargs(kwargs);
Expand Down Expand Up @@ -155,16 +157,16 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
void Probs(DataView<double, 1> &probs) override;
void PartialProbs(DataView<double, 1> &probs,
const std::vector<QubitIdType> &wires) override;
void Sample(DataView<double, 2> &samples, std::size_t shots) override;
void Sample(DataView<double, 2> &samples, const std::size_t shots) override;
void PartialSample(DataView<double, 2> &samples,
const std::vector<QubitIdType> &wires,
std::size_t shots) override;
const std::size_t shots) override;
void Counts(DataView<double, 1> &eigvals, DataView<int64_t, 1> &counts,
std::size_t shots) override;
const std::size_t shots) override;
void PartialCounts(DataView<double, 1> &eigvals,
DataView<int64_t, 1> &counts,
const std::vector<QubitIdType> &wires,
std::size_t shots) override;
const std::size_t shots) override;
mlxd marked this conversation as resolved.
Show resolved Hide resolved
auto Measure(QubitIdType wire,
std::optional<int32_t> postselect = std::nullopt)
-> Result override;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1777,3 +1777,46 @@ TEST_CASE("Measurement with a seeded device", "[Measures]") {
CHECK(*m == *m1);
}
}

TEST_CASE("Sample with a seeded device", "[Measures]") {
mlxd marked this conversation as resolved.
Show resolved Hide resolved
for (std::size_t _ = 0; _ < 5; _++) {
paul0403 marked this conversation as resolved.
Show resolved Hide resolved
std::unique_ptr<LKSimulator> sim = std::make_unique<LKSimulator>();
std::unique_ptr<LKSimulator> sim1 = std::make_unique<LKSimulator>();
mlxd marked this conversation as resolved.
Show resolved Hide resolved

std::size_t shots = 100;
mlxd marked this conversation as resolved.
Show resolved Hide resolved

std::vector<double> samples(shots * 4);
mlxd marked this conversation as resolved.
Show resolved Hide resolved
MemRefT<double, 2> buffer{
samples.data(), samples.data(), 0, {shots, 1}, {1, 1}};
DataView<double, 2> view(buffer.data_aligned, buffer.offset,
buffer.sizes, buffer.strides);

std::vector<double> samples1(shots * 4);
MemRefT<double, 2> buffer1{
samples1.data(), samples1.data(), 0, {shots, 1}, {1, 1}};
DataView<double, 2> view1(buffer1.data_aligned, buffer1.offset,
buffer1.sizes, buffer1.strides);
mlxd marked this conversation as resolved.
Show resolved Hide resolved

std::mt19937 gen(37);
mlxd marked this conversation as resolved.
Show resolved Hide resolved
sim->SetDevicePRNG(&gen);
std::vector<intptr_t> Qs;
Qs.reserve(1);
Qs.push_back(sim->AllocateQubit());
sim->NamedOperation("Hadamard", {}, {Qs[0]}, false);
sim->NamedOperation("RX", {0.5}, {Qs[0]}, false);
sim->Sample(view, shots);
mlxd marked this conversation as resolved.
Show resolved Hide resolved

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);
sim1->NamedOperation("RX", {0.5}, {Qs1[0]}, false);
sim1->Sample(view1, shots);

for (std::size_t i = 0; i < shots * 4; i++) {
CHECK((samples[i] == samples1[i]));
}
mlxd marked this conversation as resolved.
Show resolved Hide resolved
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
#pragma once
#include <chrono>
#include <cstdint>
#include <optional>

#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>
Expand Down Expand Up @@ -649,13 +650,16 @@ 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) -> std::vector<std::size_t> {
auto generate_samples(std::size_t num_samples,
paul0403 marked this conversation as resolved.
Show resolved Hide resolved
const std::optional<std::size_t> &seed = std::nullopt)
paul0403 marked this conversation as resolved.
Show resolved Hide resolved
mlxd marked this conversation as resolved.
Show resolved Hide resolved
-> 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",
Expand All @@ -674,10 +678,12 @@ class Measurements final
});

// Sampling using Random_XorShift64_Pool
Kokkos::Random_XorShift64_Pool<> rand_pool(
std::chrono::high_resolution_clock::now()
.time_since_epoch()
.count());
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());

Kokkos::parallel_for(
Kokkos::RangePolicy<KokkosExecSpace>(0, num_samples),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
#include <algorithm>
#include <complex>
#include <cstdio>
#include <optional>
#include <random>
#include <type_traits>
#include <unordered_map>
Expand Down Expand Up @@ -573,30 +574,39 @@ class Measurements final
* Reference: https://en.wikipedia.org/wiki/Alias_method
*
* @param num_samples The number of samples to generate.
* @param seed Seed to generate the samples from
* @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::size_t> &seed = std::nullopt) {
mlxd marked this conversation as resolved.
Show resolved Hide resolved
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, seed);
}

/**
* @brief Generate samples.
*
* @param wires Sample are generated for the specified wires.
* @param num_samples The number of samples to generate.
* @param seed Seed to generate the samples from
* @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) {
const std::size_t num_samples,
const std::optional<std::size_t> &seed = std::nullopt) {
mlxd marked this conversation as resolved.
Show resolved Hide resolved
const std::size_t n_wires = wires.size();
std::vector<std::size_t> samples(num_samples * n_wires);
this->setRandomSeed();
if (seed.has_value()) {
this->setSeed(seed.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 -
Expand Down
Loading