Fix dynamic wires with samples #1321

rniczh · 2026-01-06T20:51:45Z

Context:

from catalyst import *
import pennylane as qml

qml.capture.enable()

@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])

print(circuit())

It get unexpected result due to the aux wires involved.

[[0 0]
 [0 1]
 [0 0]
 [1 0]
 [0 0]
 [0 0]
 [0 1]
 [0 0]
 [1 0]
 [0 0]]

Description of the Change:

~~Only take the device wires for generating samples~~
Compact state vector before measurement.

For state vector normalization:

$|\psi| = \sqrt{\sum_i |\psi_i|^2}$
$normalized_i = \psi_i / |\psi|$

Benefits:

Possible Drawbacks:

Related GitHub Issues:

PennyLaneAI/catalyst#2339
[[sc-107206]]

github-actions · 2026-01-06T20:51:57Z

Hello. You may have forgotten to update the changelog!
Please edit .github/CHANGELOG.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

dime10

Thank you @rniczh !

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

dime10 · 2026-01-06T21:24:40Z

pennylane_lightning/core/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp

+        const std::vector<std::size_t> &wires, const std::string &kernelname,
+        std::size_t num_burnin, std::size_t num_samples) {
        std::uniform_real_distribution<PrecisionT> distrib(0.0, 1.0);
+        std::size_t num_qubits = wires.size();


Do we know that this is sufficient to identify which wires should be selected in a more complex dynamic (de)allocation scheme, or does it rely on the assumption that the selected wires are always the bottom n wires?

Mentioned this offline but I think we'd be better off providing a universal solution that we never got around to implementing (see #1254 (comment)).

I updated the PR, basically removed the original solution that only pass the device_wires to GenerateSample. Now, it will compact the state vector before the measurement process to ensure the erased qubits are erased from state vector as well.
~~ca0ff43~~ Edit: 77fbbe9

Wow that was fast! 🤩
@maliasadi are you able to look this over?

codecov · 2026-01-07T20:06:30Z

Codecov Report

❌ Patch coverage is 82.55814% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.22%. Comparing base (e7d82a2) to head (5bac5ae).

Files with missing lines	Patch %	Lines
..._lightning/core/catalyst/LightningQubitManager.hpp	0.00%	12 Missing ⚠️
...rs/lightning_qubit/catalyst/LightningSimulator.cpp	94.82%	3 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (e7d82a2) and HEAD (5bac5ae). Click for more details.

HEAD has 6 uploads less than BASE

Flag BASE (e7d82a2) HEAD (5bac5ae)

unit_tests 13 7

Additional details and impacted files

@@              Coverage Diff               @@
##           v0.44.0_rc    #1321      +/-   ##
==============================================
- Coverage       96.01%   89.22%   -6.80%     
==============================================
  Files             307      187     -120     
  Lines           46923    29012   -17911     
==============================================
- Hits            45055    25886   -19169     
- Misses           1868     3126    +1258

Flag	Coverage Δ
unit_tests	`89.22% <82.55%> (-6.80%)`	⬇️

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

maliasadi

Thanks @rniczh! A few things:

If we need this to go out with the release, you will need to rebase it against #1317 and sync ASAP with @jzaia18
Update Lighnting-Kokkos and Lightning-GPU catalyst impl similar to Lightning-Qubit
Please add a changelog entry for this fix

--------- Co-authored-by: Ali Asadi <[email protected]>

jzaia18

Nice work Hongsheng! Just a few questions and comments. ~~Is there any reason we're implementing this for lightning.qubit but not the other devices? Does this issue exist for the others as well?~~ nvm, just saw Ali already addressed this

jzaia18 · 2026-01-08T16:24:57Z

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

-    // get device wires
-    auto &&dev_wires = getDeviceWires(wires);
-
    auto li_samples = this->GenerateSamples(device_shots);

+    // Get device wires
+    auto &&dev_wires = getDeviceWires(wires);
+


Any reason for switching the order of these steps?

That's because the getMeasurements() is within the GenerateSamples. The getMeasurements() will compact the state vector to get the measurement object, that means if the getDeviceWires happens before GenerateSamples may not get the correct device wires.

std::vector<size_t> LightningSimulator::GenerateSamples(size_t shots) { auto m = getMeasurements(); // ... rest of code ...

jzaia18 · 2026-01-08T16:41:57Z

pennylane_lightning/core/catalyst/LightningQubitManager.hpp

+        std::unordered_set<DeviceQubitID> new_free_qubits;
+        for (auto old_device_id : this->free_device_qubits) {
+            auto it = old_to_new_mapping.find(old_device_id);
+            if (it != old_to_new_mapping.end()) {
+                new_free_qubits.insert(it->second);
+            } else {
+                new_free_qubits.insert(old_device_id);
+            }
+        }
+        this->free_device_qubits = std::move(new_free_qubits);


How often do we expect this section to be run? Will it be expensive to create and free the free_device_qubits set each time this is run as we currently do?

jzaia18 · 2026-01-08T16:52:27Z

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp

+    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];
+    }


This outer loop goes over the entire statevector, we probably want to use parallelism to prevent this from being a very expensive operation. Unsure if it's worth doing here or fixing in a follow-up PR. @maliasadi

This comment has been minimized.

Sign in to view

dime10 reviewed Jan 6, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

maliasadi self-requested a review January 8, 2026 15:06

maliasadi reviewed Jan 8, 2026

View reviewed changes

rniczh changed the base branch from master to v0.44.0_rc January 8, 2026 15:19

rniczh added 5 commits January 8, 2026 10:24

rebase

79451be

Fix dynamic wires with samples

7078d0e

fix

ddfdb53

add back the comment

9d19fa9

Compact state vector before measurement

7ae4e5e

jzaia18 added this to the 0.44.0 milestone Jan 8, 2026

jzaia18 and others added 8 commits January 8, 2026 10:28

Rebase

b3e806f

--------- Co-authored-by: Ali Asadi <[email protected]>

update

d365d2d

fix

6771882

no need to collapse released qubit during the

0e11617

fix compact timing

635370d

should use big-endian

aa98ec9

Add test

a8367bd

update test

6e9f177

rniczh force-pushed the rniczh/fix-dynamic-wires-with-samples branch from d3abf70 to 6e9f177 Compare January 8, 2026 15:32

rniczh changed the base branch from v0.44.0_rc to master January 8, 2026 15:37

rniczh changed the base branch from master to v0.44.0_rc January 8, 2026 15:37

clean

5bac5ae

jzaia18 reviewed Jan 8, 2026

View reviewed changes

Fix dynamic wires with samples #1321

Are you sure you want to change the base?

Fix dynamic wires with samples #1321

Conversation

rniczh commented Jan 6, 2026 • edited by sengthai Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jan 6, 2026

Uh oh!

This comment has been minimized.

dime10 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dime10 Jan 6, 2026

Choose a reason for hiding this comment

Uh oh!

dime10 Jan 6, 2026

Choose a reason for hiding this comment

Uh oh!

rniczh Jan 7, 2026 • edited by dime10 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dime10 Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

codecov bot commented Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

This comment has been minimized.

maliasadi left a comment

Choose a reason for hiding this comment

Uh oh!

jzaia18 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jzaia18 Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

rniczh Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

jzaia18 Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

jzaia18 Jan 8, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

rniczh commented Jan 6, 2026 •

edited by sengthai

Loading

rniczh Jan 7, 2026 •

edited by dime10

Loading

codecov bot commented Jan 7, 2026 •

edited

Loading

jzaia18 left a comment •

edited

Loading