diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 4f6bd2e6f8..a56120ae73 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -34,6 +34,9 @@
Bug fixes 🐛
+- Fixed sampling with dynamically allocated and released qubits in Catalyst. The state vector is now reduced before measurements to correctly handle released qubits, ensuring `qml.sample()` and other measurement operations return correct results when using `qml.allocate()` and `qml.release()`.
+ [(#1321)](https://github.com/PennyLaneAI/pennylane-lightning/pull/1321)
+
- Corrected an issue in tests where a PennyLane operator was used within a QNode to compute a
matrix, which would lead to wrongful queuing as of PennyLane
pull request [#8131](https://github.com/PennyLaneAI/pennylane/pull/8131).
diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
index 9a7c11c15f..a108f5f6ca 100644
--- a/pennylane_lightning/core/_version.py
+++ b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
Version number (major.minor.patch[-label])
"""
-__version__ = "0.44.0-rc2"
+__version__ = "0.44.0-rc3"
diff --git a/pennylane_lightning/core/catalyst/LightningQubitManager.hpp b/pennylane_lightning/core/catalyst/LightningQubitManager.hpp
index 182f376399..111ada288b 100644
--- a/pennylane_lightning/core/catalyst/LightningQubitManager.hpp
+++ b/pennylane_lightning/core/catalyst/LightningQubitManager.hpp
@@ -163,5 +163,16 @@ class QubitManager final {
this->qubit_id_map.clear();
this->free_device_qubits.clear();
}
+
+ void RemapDeviceIds(const std::unordered_map
+ &old_to_new_mapping) {
+ // Update each program_id's device_id according to the mapping
+ for (auto &[program_id, device_id] : this->qubit_id_map) {
+ if (auto it = old_to_new_mapping.find(device_id);
+ it != old_to_new_mapping.end()) {
+ device_id = it->second;
+ }
+ }
+ }
};
} // namespace Catalyst::Runtime::Simulator
diff --git a/pennylane_lightning/core/catalyst/tests/Test_LightningDriver.cpp b/pennylane_lightning/core/catalyst/tests/Test_LightningDriver.cpp
index f2a1850d7b..2bbdaacd87 100644
--- a/pennylane_lightning/core/catalyst/tests/Test_LightningDriver.cpp
+++ b/pennylane_lightning/core/catalyst/tests/Test_LightningDriver.cpp
@@ -331,3 +331,85 @@ TEST_CASE("Release Qubits", "[Driver]") {
CHECK(sim->GetNumQubits() == 2);
}
+
+TEST_CASE("Sample after dynamic qubit release", "[Driver]") {
+ // This test mirrors the Python code:
+ // @qjit
+ // @qml.qnode(qml.device("lightning.qubit", wires=3, shots=10))
+ // def circuit():
+ // with qml.allocate(2) as qs:
+ // qml.X(qs[1])
+ // return qml.sample(wires=[0, 1])
+
+ std::unique_ptr sim = std::make_unique();
+
+ // Allocate 3 static qubits (wires 0, 1, 2) : all in |0>
+ std::vector static_qubits = sim->AllocateQubits(3);
+
+ // Dynamically allocate 2 qubits
+ std::vector dynamic_qubits = sim->AllocateQubits(2);
+
+ // Apply PauliX to dynamic_qubits[1]
+ sim->NamedOperation("PauliX", {}, {dynamic_qubits[1]}, false);
+
+ // Release the dynamic qubits
+ sim->ReleaseQubits(dynamic_qubits);
+
+ // Sample on static wires [0, 1]
+ // Since static qubits were never modified, they should all be |0>
+ constexpr size_t num_shots = 10;
+ constexpr size_t num_wires = 2;
+ sim->SetDeviceShots(num_shots);
+
+ std::vector samples(num_shots * num_wires);
+ const size_t sizes[2] = {num_shots, num_wires};
+ const size_t strides[2] = {num_wires, 1}; // row-major: stride[0]=num_wires
+ DataView samples_view(samples.data(), 0, sizes, strides);
+
+ sim->PartialSample(samples_view, {static_qubits[0], static_qubits[1]});
+
+ for (size_t i = 0; i < num_shots * num_wires; i++) {
+ CHECK(samples[i] == 0.);
+ }
+}
+
+TEST_CASE("Sample after releasing middle qubit (triggers remap)", "[Driver]") {
+ // Scenario:
+ // 1. Allocate 3 qubits -> device IDs: 0, 1, 2
+ // 2. Apply X to qubit 2 (device ID 2) -> state |001>
+ // 3. Release qubit 1 (device ID 1) -> remaining device IDs: 0, 2
+ // 4. reduceStateVector remaps: device ID 2 -> 1
+ // 5. Sample qubit 2 (now device ID 1) -> should get |1>
+
+ std::unique_ptr sim = std::make_unique();
+
+ // Allocate 3 qubits: device IDs 0, 1, 2
+ std::vector qubits = sim->AllocateQubits(3);
+
+ // Apply X to qubit[2] (device ID 2): state becomes |001>
+ sim->NamedOperation("PauliX", {}, {qubits[2]}, false);
+
+ // Release qubit[1] (device ID 1), this creates a gap in device IDs
+ // Remaining: qubit[0] (device 0), qubit[2] (device 2)
+ // After reduceStateVector: qubit[0] -> device 0, qubit[2] -> device 1
+ sim->ReleaseQubit(qubits[1]);
+
+ // Sample on qubit[0] and qubit[2]
+ // qubit[0] should be |0>, qubit[2] should be |1>
+ constexpr size_t num_shots = 10;
+ constexpr size_t num_wires = 2;
+ sim->SetDeviceShots(num_shots);
+
+ std::vector samples(num_shots * num_wires);
+ const size_t sizes[2] = {num_shots, num_wires};
+ const size_t strides[2] = {num_wires, 1}; // row-major: stride[0]=num_wires
+ DataView samples_view(samples.data(), 0, sizes, strides);
+
+ sim->PartialSample(samples_view, {qubits[0], qubits[2]});
+
+ // each shot should be [0, 1] (qubit[0]=0, qubit[2]=1)
+ for (size_t shot = 0; shot < num_shots; shot++) {
+ CHECK(samples[shot * num_wires + 0] == 0.); // qubit[0] is |0>
+ CHECK(samples[shot * num_wires + 1] == 1.); // qubit[2] is |1>
+ }
+}
diff --git a/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp b/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp
index c5943ac1bf..aa94839b4d 100644
--- a/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp
+++ b/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp
@@ -15,6 +15,7 @@
#include
#include "LightningGPUSimulator.hpp"
+#include "Util.hpp"
namespace Catalyst::Runtime::Simulator {
@@ -82,9 +83,14 @@ auto LightningGPUSimulator::AllocateQubits(std::size_t num_qubits)
}
void LightningGPUSimulator::ReleaseQubit(QubitIdType q) {
- // We do not deallocate physical memory in the statevector for this
- // operation, instead we just mark the qubits as released.
+ RT_FAIL_IF(!this->qubit_manager.isValidQubitId(q),
+ "Invalid qubit to release");
+
+ // Mark the qubit as released in the qubit manager
this->qubit_manager.Release(q);
+
+ // Mark that reduction is needed
+ this->needs_reduction = true;
}
void LightningGPUSimulator::ReleaseQubits(const std::vector &ids) {
@@ -101,12 +107,13 @@ void LightningGPUSimulator::ReleaseQubits(const std::vector &ids) {
if (deallocate_all) {
this->qubit_manager.ReleaseAll();
this->device_sv = std::make_unique(0);
+ this->needs_reduction = false;
return;
}
}
for (auto id : ids) {
- this->qubit_manager.Release(id);
+ this->ReleaseQubit(id);
}
}
@@ -114,6 +121,81 @@ auto LightningGPUSimulator::GetNumQubits() const -> std::size_t {
return this->qubit_manager.getNumQubits();
}
+auto LightningGPUSimulator::getMeasurements()
+ -> Pennylane::LightningGPU::Measures::Measurements {
+ reduceStateVector();
+ return Pennylane::LightningGPU::Measures::Measurements{
+ *(this->device_sv)};
+}
+
+void LightningGPUSimulator::reduceStateVector() {
+ if (!this->needs_reduction) {
+ return;
+ }
+
+ // Get active qubits
+ auto all_qubits = this->qubit_manager.getAllQubitIds();
+ std::vector> wire_id_pairs;
+
+ for (auto qid : all_qubits) {
+ size_t device_wire = this->qubit_manager.getDeviceId(qid);
+ wire_id_pairs.push_back({device_wire, qid});
+ }
+
+ // Sort by device wire index
+ std::sort(wire_id_pairs.begin(), wire_id_pairs.end());
+
+ // Extract reduced state vector - need to copy from GPU first
+ auto old_data = this->device_sv->getDataVector();
+ size_t num_qubits_after = wire_id_pairs.size();
+ size_t new_size = 1UL << num_qubits_after;
+
+ std::vector> new_data(new_size);
+
+ // state[idx] = |q0 q1 ... q_{n-1}>
+ // where q_i = (idx >> (n-1-i)) & 1
+ // So device wire 0 corresponds to the MSB (bit n-1)
+ size_t old_num_qubits = this->device_sv->getNumQubits();
+
+ for (size_t old_idx = 0; old_idx < old_data.size(); old_idx++) {
+ size_t new_idx = 0;
+ for (size_t i = 0; i < num_qubits_after; i++) {
+ size_t old_wire = wire_id_pairs[i].first;
+ size_t old_bit_pos = old_num_qubits - 1 - old_wire;
+ size_t new_bit_pos = num_qubits_after - 1 - i;
+
+ if ((old_idx >> old_bit_pos) & 1) {
+ new_idx |= (1UL << new_bit_pos);
+ }
+ }
+
+ new_data[new_idx] += old_data[old_idx];
+ }
+
+ // Normalize the state vector
+ Pennylane::Util::normalizeStateVector(new_data);
+
+ // Replace the state vector
+ this->device_sv =
+ std::make_unique(new_data.data(), new_data.size());
+
+ // Remap device ids
+ std::unordered_map old_to_new_device_id;
+ for (size_t new_idx = 0; new_idx < wire_id_pairs.size(); new_idx++) {
+ size_t old_device_id = wire_id_pairs[new_idx].first;
+ size_t new_device_id = new_idx;
+ if (old_device_id != new_device_id) {
+ old_to_new_device_id[old_device_id] = new_device_id;
+ }
+ }
+
+ if (!old_to_new_device_id.empty()) {
+ this->qubit_manager.RemapDeviceIds(old_to_new_device_id);
+ }
+
+ this->needs_reduction = false;
+}
+
void LightningGPUSimulator::StartTapeRecording() {
RT_FAIL_IF(this->tape_recording, "Cannot re-activate the cache manager");
this->tape_recording = true;
@@ -264,9 +346,7 @@ auto LightningGPUSimulator::Expval(ObsIdType obsKey) -> double {
auto &&obs = this->obs_manager.getObservable(obsKey);
- Pennylane::LightningGPU::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
return device_shots ? m.expval(*obs, device_shots, {}) : m.expval(*obs);
@@ -283,9 +363,7 @@ auto LightningGPUSimulator::Var(ObsIdType obsKey) -> double {
auto &&obs = this->obs_manager.getObservable(obsKey);
- Pennylane::LightningGPU::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
return device_shots ? m.var(*obs, device_shots) : m.var(*obs);
@@ -307,9 +385,7 @@ void LightningGPUSimulator::State(DataView, 1> &state) {
}
void LightningGPUSimulator::Probs(DataView &probs) {
- Pennylane::LightningGPU::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
auto &&dv_probs = device_shots ? m.probs(device_shots) : m.probs();
@@ -328,12 +404,11 @@ void LightningGPUSimulator::PartialProbs(
RT_FAIL_IF(numWires > numQubits, "Invalid number of wires");
RT_FAIL_IF(!isValidQubits(wires), "Invalid given wires to measure");
- auto dev_wires = getDeviceWires(wires);
- Pennylane::LightningGPU::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
+ auto dev_wires = getDeviceWires(wires);
+
auto &&dv_probs =
device_shots ? m.probs(dev_wires, device_shots) : m.probs(dev_wires);
@@ -344,10 +419,7 @@ void LightningGPUSimulator::PartialProbs(
}
std::vector LightningGPUSimulator::GenerateSamples(size_t shots) {
- // generate_samples is a member function of the Measures class.
- Pennylane::LightningGPU::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
return m.generate_samples(shots);
@@ -383,11 +455,11 @@ void LightningGPUSimulator::PartialSample(
RT_FAIL_IF(samples.size() != device_shots * numWires,
"Invalid size for the pre-allocated partial-samples");
- // get device wires
- auto &&dev_wires = getDeviceWires(wires);
-
auto li_samples = this->GenerateSamples(device_shots);
+ // Get device wires
+ auto &&dev_wires = getDeviceWires(wires);
+
// The lightning samples are layed out as a single vector of size
// shots*qubits, where each element represents a single bit. The
// corresponding shape is (shots, qubits). Gather the desired bits
@@ -444,11 +516,11 @@ void LightningGPUSimulator::PartialCounts(
RT_FAIL_IF((eigvals.size() != numElements || counts.size() != numElements),
"Invalid size for the pre-allocated partial-counts");
- // get device wires
- auto &&dev_wires = getDeviceWires(wires);
-
auto li_samples = this->GenerateSamples(device_shots);
+ // Get device wires
+ auto &&dev_wires = getDeviceWires(wires);
+
// Fill the eigenvalues with the integer representation of the
// corresponding computational basis bitstring. In the future,
// eigenvalues can also be obtained from an observable, hence the
diff --git a/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.hpp b/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.hpp
index 947ba9abf0..e8bba0c442 100644
--- a/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.hpp
+++ b/pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.hpp
@@ -67,6 +67,9 @@ class LightningGPUSimulator final : public Catalyst::Runtime::QuantumDevice {
std::unique_ptr device_sv = std::make_unique(0);
LightningGPUObsManager obs_manager{};
+ // Flag to indicate if state vector needs reduction
+ bool needs_reduction{false};
+
inline auto isValidQubit(QubitIdType wire) -> bool {
return this->qubit_manager.isValidQubitId(wire);
}
@@ -104,6 +107,13 @@ class LightningGPUSimulator final : public Catalyst::Runtime::QuantumDevice {
auto GenerateSamples(size_t shots) -> std::vector;
+ // Reduce state vector by removing released qubits
+ void reduceStateVector();
+
+ // Helper to get Measurements object with reduced state vector
+ auto getMeasurements()
+ -> Pennylane::LightningGPU::Measures::Measurements;
+
public:
explicit LightningGPUSimulator(const std::string &kwargs = "{}") {
auto &&args = Catalyst::Runtime::parse_kwargs(kwargs);
diff --git a/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp b/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp
index 82f550d3ee..11196631f0 100644
--- a/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp
+++ b/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.cpp
@@ -18,6 +18,7 @@
#include
#include "LightningKokkosSimulator.hpp"
+#include "Util.hpp"
namespace Catalyst::Runtime::Simulator {
@@ -84,9 +85,14 @@ auto LightningKokkosSimulator::AllocateQubits(std::size_t num_qubits)
}
void LightningKokkosSimulator::ReleaseQubit(QubitIdType q) {
- // We do not deallocate physical memory in the statevector for this
- // operation, instead we just mark the qubits as released.
+ RT_FAIL_IF(!this->qubit_manager.isValidQubitId(q),
+ "Invalid qubit to release");
+
+ // Mark the qubit as released in the qubit manager
this->qubit_manager.Release(q);
+
+ // Mark that reduction is needed
+ this->needs_reduction = true;
}
void LightningKokkosSimulator::ReleaseQubits(
@@ -104,12 +110,13 @@ void LightningKokkosSimulator::ReleaseQubits(
if (deallocate_all) {
this->qubit_manager.ReleaseAll();
this->device_sv = std::make_unique(0);
+ this->needs_reduction = false;
return;
}
}
for (auto id : ids) {
- this->qubit_manager.Release(id);
+ this->ReleaseQubit(id);
}
}
@@ -117,6 +124,80 @@ auto LightningKokkosSimulator::GetNumQubits() const -> std::size_t {
return this->qubit_manager.getNumQubits();
}
+auto LightningKokkosSimulator::getMeasurements()
+ -> Pennylane::LightningKokkos::Measures::Measurements {
+ reduceStateVector();
+ return Pennylane::LightningKokkos::Measures::Measurements{
+ *(this->device_sv)};
+}
+
+void LightningKokkosSimulator::reduceStateVector() {
+ if (!this->needs_reduction) {
+ return;
+ }
+
+ // Get active qubits
+ auto all_qubits = this->qubit_manager.getAllQubitIds();
+ std::vector> wire_id_pairs;
+
+ for (auto qid : all_qubits) {
+ size_t device_wire = this->qubit_manager.getDeviceId(qid);
+ wire_id_pairs.push_back({device_wire, qid});
+ }
+
+ // Sort by device wire index
+ std::sort(wire_id_pairs.begin(), wire_id_pairs.end());
+
+ // Extract reduced state vector
+ auto old_data = this->device_sv->getDataVector();
+ size_t num_qubits_after = wire_id_pairs.size();
+ size_t new_size = 1UL << num_qubits_after;
+
+ std::vector> new_data(new_size);
+
+ // state[idx] = |q0 q1 ... q_{n-1}>
+ // where q_i = (idx >> (n-1-i)) & 1
+ // So device wire 0 corresponds to the MSB (bit n-1)
+ size_t old_num_qubits = this->device_sv->getNumQubits();
+
+ for (size_t old_idx = 0; old_idx < old_data.size(); old_idx++) {
+ size_t new_idx = 0;
+ for (size_t i = 0; i < num_qubits_after; i++) {
+ size_t old_wire = wire_id_pairs[i].first;
+ size_t old_bit_pos = old_num_qubits - 1 - old_wire;
+ size_t new_bit_pos = num_qubits_after - 1 - i;
+
+ if ((old_idx >> old_bit_pos) & 1) {
+ new_idx |= (1UL << new_bit_pos);
+ }
+ }
+
+ new_data[new_idx] += old_data[old_idx];
+ }
+
+ // Normalize the state vector
+ Pennylane::Util::normalizeStateVector(new_data);
+
+ // Replace the state vector
+ this->device_sv = std::make_unique(new_data);
+
+ // Remap device ids
+ std::unordered_map old_to_new_device_id;
+ for (size_t new_idx = 0; new_idx < wire_id_pairs.size(); new_idx++) {
+ size_t old_device_id = wire_id_pairs[new_idx].first;
+ size_t new_device_id = new_idx;
+ if (old_device_id != new_device_id) {
+ old_to_new_device_id[old_device_id] = new_device_id;
+ }
+ }
+
+ if (!old_to_new_device_id.empty()) {
+ this->qubit_manager.RemapDeviceIds(old_to_new_device_id);
+ }
+
+ this->needs_reduction = false;
+}
+
void LightningKokkosSimulator::StartTapeRecording() {
RT_FAIL_IF(this->tape_recording, "Cannot re-activate the cache manager");
this->tape_recording = true;
@@ -282,9 +363,7 @@ auto LightningKokkosSimulator::Expval(ObsIdType obsKey) -> double {
auto &&obs = this->obs_manager.getObservable(obsKey);
- Pennylane::LightningKokkos::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
return device_shots ? m.expval(*obs, device_shots, {}) : m.expval(*obs);
@@ -301,9 +380,7 @@ auto LightningKokkosSimulator::Var(ObsIdType obsKey) -> double {
auto &&obs = this->obs_manager.getObservable(obsKey);
- Pennylane::LightningKokkos::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
return device_shots ? m.var(*obs, device_shots) : m.var(*obs);
@@ -333,9 +410,7 @@ void LightningKokkosSimulator::State(DataView, 1> &state) {
}
void LightningKokkosSimulator::Probs(DataView &probs) {
- Pennylane::LightningKokkos::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
auto &&dv_probs = device_shots ? m.probs(device_shots) : m.probs();
@@ -354,12 +429,11 @@ void LightningKokkosSimulator::PartialProbs(
RT_FAIL_IF(numWires > numQubits, "Invalid number of wires");
RT_FAIL_IF(!isValidQubits(wires), "Invalid given wires to measure");
- auto dev_wires = getDeviceWires(wires);
- Pennylane::LightningKokkos::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
+ auto dev_wires = getDeviceWires(wires);
+
auto &&dv_probs =
device_shots ? m.probs(dev_wires, device_shots) : m.probs(dev_wires);
@@ -370,10 +444,7 @@ void LightningKokkosSimulator::PartialProbs(
}
std::vector LightningKokkosSimulator::GenerateSamples(size_t shots) {
- // generate_samples is a member function of the Measures class.
- Pennylane::LightningKokkos::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
// PL-Lightning-Kokkos generates samples using the alias method.
@@ -411,11 +482,11 @@ void LightningKokkosSimulator::PartialSample(
RT_FAIL_IF(samples.size() != device_shots * numWires,
"Invalid size for the pre-allocated partial-samples");
- // get device wires
- auto &&dev_wires = getDeviceWires(wires);
-
auto li_samples = this->GenerateSamples(device_shots);
+ // Get device wires
+ auto &&dev_wires = getDeviceWires(wires);
+
// The lightning samples are layed out as a single vector of size
// shots*qubits, where each element represents a single bit. The
// corresponding shape is (shots, qubits). Gather the desired bits
@@ -472,11 +543,11 @@ void LightningKokkosSimulator::PartialCounts(
RT_FAIL_IF((eigvals.size() != numElements || counts.size() != numElements),
"Invalid size for the pre-allocated partial-counts");
- // get device wires
- auto &&dev_wires = getDeviceWires(wires);
-
auto li_samples = this->GenerateSamples(device_shots);
+ // Get device wires
+ auto &&dev_wires = getDeviceWires(wires);
+
// Fill the eigenvalues with the integer representation of the
// corresponding computational basis bitstring. In the future,
// eigenvalues can also be obtained from an observable, hence the
diff --git a/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.hpp b/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.hpp
index 10c4921eee..4ce3118af9 100644
--- a/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.hpp
+++ b/pennylane_lightning/core/simulators/lightning_kokkos/catalyst/LightningKokkosSimulator.hpp
@@ -67,6 +67,9 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
std::unique_ptr device_sv = std::make_unique(0);
LightningKokkosObsManager obs_manager{};
+ // Flag to indicate if state vector needs reduction
+ bool needs_reduction{false};
+
inline auto isValidQubit(QubitIdType wire) -> bool {
return this->qubit_manager.isValidQubitId(wire);
}
@@ -104,6 +107,13 @@ class LightningKokkosSimulator final : public Catalyst::Runtime::QuantumDevice {
auto GenerateSamples(size_t shots) -> std::vector;
+ // Reduce state vector by removing released qubits
+ void reduceStateVector();
+
+ // Helper to get Measurements object with reduced state vector
+ auto getMeasurements()
+ -> Pennylane::LightningKokkos::Measures::Measurements;
+
public:
explicit LightningKokkosSimulator(
const std::string &kwargs = "{}") noexcept {
diff --git a/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp b/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp
index fd68d1b205..9e9932e18f 100644
--- a/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp
+++ b/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp
@@ -20,6 +20,7 @@
#include "JacobianData.hpp"
#include "LinearAlgebra.hpp"
#include "MeasurementsLQubit.hpp"
+#include "Util.hpp"
namespace Catalyst::Runtime::Simulator {
@@ -101,12 +102,13 @@ void LightningSimulator::ReleaseQubits(const std::vector &ids) {
if (deallocate_all) {
this->qubit_manager.ReleaseAll();
this->device_sv = std::make_unique(0);
+ this->needs_reduction = false;
return;
}
}
for (auto id : ids) {
- this->qubit_manager.Release(id);
+ this->ReleaseQubit(id);
}
}
@@ -117,10 +119,88 @@ void LightningSimulator::ReleaseQubit(QubitIdType q) {
// maintained,
// in particular: measurements must not compute results for released
// qubits, only for active/allocated qubits
+ RT_FAIL_IF(!this->qubit_manager.isValidQubitId(q),
+ "Invalid qubit to release");
+
+ // Mark the qubit as released in the qubit manager
this->qubit_manager.Release(q);
- // TODO: We don't have to check whether the released qubit is pure, but it
- // is something we can add easily to catch program errors.
+ // Mark that reduction is needed
+ this->needs_reduction = true;
+}
+
+auto LightningSimulator::getMeasurements()
+ -> Pennylane::LightningQubit::Measures::Measurements {
+ reduceStateVector();
+ return Pennylane::LightningQubit::Measures::Measurements{
+ *(this->device_sv)};
+}
+
+void LightningSimulator::reduceStateVector() {
+ if (!this->needs_reduction) {
+ return;
+ }
+
+ // Get active qubits
+ auto all_qubits = this->qubit_manager.getAllQubitIds();
+ std::vector> wire_id_pairs;
+
+ for (auto qid : all_qubits) {
+ size_t device_wire = this->qubit_manager.getDeviceId(qid);
+ wire_id_pairs.push_back({device_wire, qid});
+ }
+
+ // Sort by device wire index
+ std::sort(wire_id_pairs.begin(), wire_id_pairs.end());
+
+ // Extract reduced state vector
+ auto &old_data = this->device_sv->getDataVector();
+ size_t num_qubits_after = wire_id_pairs.size();
+ size_t new_size = 1UL << num_qubits_after;
+
+ std::vector> new_data(new_size);
+
+ // state[idx] = |q0 q1 ... q_{n-1}>
+ // where q_i = (idx >> (n-1-i)) & 1
+ // So device wire 0 corresponds to the MSB (bit n-1)
+ size_t old_num_qubits = this->device_sv->getNumQubits();
+
+ for (size_t old_idx = 0; old_idx < old_data.size(); old_idx++) {
+ size_t new_idx = 0;
+ for (size_t i = 0; i < num_qubits_after; i++) {
+ size_t old_wire = wire_id_pairs[i].first;
+ size_t old_bit_pos = old_num_qubits - 1 - old_wire;
+ size_t new_bit_pos = num_qubits_after - 1 - i;
+
+ if ((old_idx >> old_bit_pos) & 1) {
+ new_idx |= (1UL << new_bit_pos);
+ }
+ }
+
+ new_data[new_idx] += old_data[old_idx];
+ }
+
+ // Normalize the state vector
+ Pennylane::Util::normalizeStateVector(new_data);
+
+ // Replace the state vector
+ this->device_sv = std::make_unique(new_data);
+
+ // Remap device ids
+ std::unordered_map old_to_new_device_id;
+ for (size_t new_idx = 0; new_idx < wire_id_pairs.size(); new_idx++) {
+ size_t old_device_id = wire_id_pairs[new_idx].first;
+ size_t new_device_id = new_idx;
+ if (old_device_id != new_device_id) {
+ old_to_new_device_id[old_device_id] = new_device_id;
+ }
+ }
+
+ if (!old_to_new_device_id.empty()) {
+ this->qubit_manager.RemapDeviceIds(old_to_new_device_id);
+ }
+
+ this->needs_reduction = false;
}
auto LightningSimulator::GetNumQubits() const -> size_t {
@@ -275,9 +355,7 @@ auto LightningSimulator::Expval(ObsIdType obsKey) -> double {
this->cache_manager.addObservable(obsKey, MeasurementsT::Expval);
}
- Pennylane::LightningQubit::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
return (device_shots != 0U) ? m.expval(*obs, device_shots, {})
@@ -294,9 +372,7 @@ auto LightningSimulator::Var(ObsIdType obsKey) -> double {
this->cache_manager.addObservable(obsKey, MeasurementsT::Var);
}
- Pennylane::LightningQubit::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
return (device_shots != 0U) ? m.var(*obs, device_shots) : m.var(*obs);
@@ -311,9 +387,7 @@ void LightningSimulator::State(DataView, 1> &state) {
}
void LightningSimulator::Probs(DataView &probs) {
- Pennylane::LightningQubit::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
auto &&dv_probs = (device_shots != 0U) ? m.probs(device_shots) : m.probs();
@@ -332,12 +406,11 @@ void LightningSimulator::PartialProbs(DataView &probs,
RT_FAIL_IF(numWires > numQubits, "Invalid number of wires");
RT_FAIL_IF(!isValidQubits(wires), "Invalid given wires to measure");
- auto dev_wires = getDeviceWires(wires);
- Pennylane::LightningQubit::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
+ auto dev_wires = getDeviceWires(wires);
+
auto &&dv_probs = (device_shots != 0U) ? m.probs(dev_wires, device_shots)
: m.probs(dev_wires);
@@ -349,10 +422,7 @@ void LightningSimulator::PartialProbs(DataView &probs,
std::vector
LightningSimulator::GenerateSamplesMetropolis(size_t shots) {
- // generate_samples_metropolis is a member function of the Measures class.
- Pennylane::LightningQubit::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
// PL-Lightning generates samples using the alias method.
@@ -367,13 +437,7 @@ LightningSimulator::GenerateSamplesMetropolis(size_t shots) {
}
std::vector LightningSimulator::GenerateSamples(size_t shots) {
- if (this->mcmc) {
- return this->GenerateSamplesMetropolis(shots);
- }
- // generate_samples is a member function of the Measures class.
- Pennylane::LightningQubit::Measures::Measurements m{
- *(this->device_sv)};
-
+ auto m = getMeasurements();
m.setSeed(this->generateSeed());
// PL-Lightning generates samples using the alias method.
@@ -417,11 +481,11 @@ void LightningSimulator::PartialSample(DataView &samples,
RT_FAIL_IF(samples.size() != device_shots * numWires,
"Invalid size for the pre-allocated partial-samples");
- // get device wires
- auto &&dev_wires = getDeviceWires(wires);
-
auto li_samples = this->GenerateSamples(device_shots);
+ // Get device wires
+ auto &&dev_wires = getDeviceWires(wires);
+
// The lightning samples are layed out as a single vector of size
// shots*qubits, where each element represents a single bit. The
// corresponding shape is (shots, qubits). Gather the desired bits
@@ -478,11 +542,11 @@ void LightningSimulator::PartialCounts(DataView &eigvals,
RT_FAIL_IF((eigvals.size() != numElements || counts.size() != numElements),
"Invalid size for the pre-allocated partial-counts");
- // get device wires
- auto &&dev_wires = getDeviceWires(wires);
-
auto li_samples = this->GenerateSamples(device_shots);
+ // Get device wires
+ auto &&dev_wires = getDeviceWires(wires);
+
// Fill the eigenvalues with the integer representation of the corresponding
// computational basis bitstring. In the future, eigenvalues can also be
// obtained from an observable, hence the bitstring integer is stored as a
diff --git a/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.hpp b/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.hpp
index 2984da8606..d2954fb39d 100644
--- a/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.hpp
+++ b/pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.hpp
@@ -23,6 +23,7 @@
#include
#include
+#include "MeasurementsLQubit.hpp"
#include "StateVectorLQubitManaged.hpp"
#include "CacheManager.hpp"
@@ -60,6 +61,9 @@ class LightningSimulator final : public Catalyst::Runtime::QuantumDevice {
std::unique_ptr device_sv = std::make_unique(0);
LightningObsManager obs_manager{};
+ // Flag to indicate if state vector needs reduction
+ bool needs_reduction{false};
+
inline auto isValidQubit(QubitIdType wire) -> bool {
return this->qubit_manager.isValidQubitId(wire);
}
@@ -98,6 +102,13 @@ class LightningSimulator final : public Catalyst::Runtime::QuantumDevice {
auto GenerateSamplesMetropolis(size_t shots) -> std::vector;
auto GenerateSamples(size_t shots) -> std::vector;
+ // Reduce state vector by removing released qubits
+ void reduceStateVector();
+
+ // Helper to get Measurements object with reduced state vector
+ auto getMeasurements()
+ -> Pennylane::LightningQubit::Measures::Measurements;
+
public:
explicit LightningSimulator(
const std::string &kwargs = "{}") // NOLINT(hicpp-member-init)
diff --git a/pennylane_lightning/core/utils/Util.hpp b/pennylane_lightning/core/utils/Util.hpp
index 6fc6f60567..79b7cdff1a 100644
--- a/pennylane_lightning/core/utils/Util.hpp
+++ b/pennylane_lightning/core/utils/Util.hpp
@@ -690,4 +690,29 @@ inline auto PL_reinterpret_cast(SrcType *src_ptr) -> DestType * {
return reinterpret_cast(src_ptr);
}
+/**
+ * @brief Normalize a complex state vector to unit norm.
+ *
+ * This function normalizes a state vector in-place
+ *
+ * @tparam ComplexT Complex number type (e.g., std::complex).
+ * @tparam Alloc Allocator type.
+ * @param data State vector data to normalize.
+ */
+template >
+inline void normalizeStateVector(std::vector &data) {
+ using PrecisionT = std::decay_t().real())>;
+ PrecisionT norm_squared = 0;
+ for (const auto &litude : data) {
+ norm_squared += amplitude.real() * amplitude.real() +
+ amplitude.imag() * amplitude.imag();
+ }
+ auto norm = std::sqrt(norm_squared);
+ if (norm > std::numeric_limits::epsilon()) {
+ for (auto &elem : data) {
+ elem /= norm;
+ }
+ }
+}
+
} // namespace Pennylane::Util
diff --git a/pennylane_lightning/core/utils/tests/Test_Util.cpp b/pennylane_lightning/core/utils/tests/Test_Util.cpp
index 1e1bdd547f..dc6c976dd4 100644
--- a/pennylane_lightning/core/utils/tests/Test_Util.cpp
+++ b/pennylane_lightning/core/utils/tests/Test_Util.cpp
@@ -183,6 +183,78 @@ TEMPLATE_TEST_CASE("Util::squaredNorm", "[Util][LinearAlgebra]", float,
}
}
+TEMPLATE_TEST_CASE("Util::normalizeStateVector", "[Util][LinearAlgebra]", float,
+ double) {
+ SECTION("Normalizing an already normalized vector") {
+ // 1/sqrt(2) * (|0> + |1>)
+ const TestType inv_sqrt2 = TestType{1} / std::sqrt(TestType{2});
+ std::vector> vec{{inv_sqrt2, 0}, {inv_sqrt2, 0}};
+
+ normalizeStateVector(vec);
+
+ // Should remain unchanged with tolerance
+ CHECK(vec[0].real() == Approx(inv_sqrt2).margin(1e-7));
+ CHECK(vec[0].imag() == Approx(0).margin(1e-7));
+ CHECK(vec[1].real() == Approx(inv_sqrt2).margin(1e-7));
+ CHECK(vec[1].imag() == Approx(0).margin(1e-7));
+ CHECK(squaredNorm(vec) == Approx(1.0));
+ }
+
+ SECTION("Normalizing an unnormalized vector") {
+ // Unnormalized vector: (1, 1)
+ // After normalization: (1/sqrt(2), 1/sqrt(2))
+ std::vector> vec{{1, 0}, {1, 0}};
+
+ normalizeStateVector(vec);
+
+ const TestType inv_sqrt2 = TestType{1} / std::sqrt(TestType{2});
+ CHECK(vec[0].real() == Approx(inv_sqrt2).margin(1e-7));
+ CHECK(vec[0].imag() == Approx(0).margin(1e-7));
+ CHECK(vec[1].real() == Approx(inv_sqrt2).margin(1e-7));
+ CHECK(vec[1].imag() == Approx(0).margin(1e-7));
+ CHECK(squaredNorm(vec) == Approx(1.0));
+ }
+
+ SECTION("Normalizing a vector with complex amplitudes") {
+ // Vector: (1+i, 1-i)
+ // squared norm = 2 + 2 = 4, norm = 2
+ std::vector> vec{{1, 1}, {1, -1}};
+
+ normalizeStateVector(vec);
+
+ CHECK(vec[0].real() == Approx(0.5).margin(1e-7));
+ CHECK(vec[0].imag() == Approx(0.5).margin(1e-7));
+ CHECK(vec[1].real() == Approx(0.5).margin(1e-7));
+ CHECK(vec[1].imag() == Approx(-0.5).margin(1e-7));
+ CHECK(squaredNorm(vec) == Approx(1.0));
+ }
+
+ SECTION("Normalizing a zero vector (edge case)") {
+ std::vector> vec{{0, 0}, {0, 0}};
+ // remain zero
+ normalizeStateVector(vec);
+
+ CHECK(vec[0].real() == Approx(0).margin(1e-7));
+ CHECK(vec[0].imag() == Approx(0).margin(1e-7));
+ CHECK(vec[1].real() == Approx(0).margin(1e-7));
+ CHECK(vec[1].imag() == Approx(0).margin(1e-7));
+ }
+
+ SECTION("Normalizing a larger state vector") {
+ // 2 qubits: (2, 0, 0, 0) -> (1, 0, 0, 0)
+ std::vector> vec{{2, 0}, {0, 0}, {0, 0}, {0, 0}};
+
+ normalizeStateVector(vec);
+
+ CHECK(vec[0].real() == Approx(1.0).margin(1e-7));
+ CHECK(vec[0].imag() == Approx(0).margin(1e-7));
+ CHECK(vec[1].real() == Approx(0).margin(1e-7));
+ CHECK(vec[2].real() == Approx(0).margin(1e-7));
+ CHECK(vec[3].real() == Approx(0).margin(1e-7));
+ CHECK(squaredNorm(vec) == Approx(1.0));
+ }
+}
+
TEMPLATE_TEST_CASE("Util::is_Hermitian", "[Util][LinearAlgebra]", float,
double) {
SECTION("Test for a Hermitian matrix") {