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

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

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

Codecov Report

❌ Patch coverage is 99.58333% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 93.89%. Comparing base (e7d82a2) to head (6c8cdb2).

Files with missing lines	Patch %	Lines
...rs/lightning_qubit/catalyst/LightningSimulator.cpp	98.00%	1 Missing ⚠️

Additional details and impacted files

@@              Coverage Diff               @@
##           v0.44.0_rc    #1321      +/-   ##
==============================================
- Coverage       96.01%   93.89%   -2.13%     
==============================================
  Files             307      243      -64     
  Lines           46923    40264    -6659     
==============================================
- Hits            45055    37804    -7251     
- Misses           1868     2460     +592

Flag	Coverage Δ
unit_tests	`93.89% <99.58%> (-2.13%)`	⬇️

☔ 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

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

pennylane_lightning/core/catalyst/LightningQubitManager.hpp

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

.github/CHANGELOG.md

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

maliasadi · 2026-01-08T22:23:19Z

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

+        *(this->device_sv)};
+}
+
+void LightningGPUSimulator::CompactStateVector() {


Where are these tested? I couldn't find any tests for these methods for all devices.

The function is triggered by getMeasurements(), ensuring that any test executing the measurement process also updates the state vector. Currently, only one test case is added (the one that reported from the issue PennyLaneAI/catalyst#2339). The specific test can be found here:

https://github.com/PennyLaneAI/pennylane-lightning/pull/1321/changes#diff-6f0e99f87b0a9753689936392636b52eb19b81aac144b24d12de297f3c617447R335

Will this be added to Catalyst for the release?

per our offline discussions, what if these aux wires are entangled, we should test and add some handling mechanism for these cases in both Lightning & Catalyst?

pyproject.toml

pennylane_lightning/core/catalyst/LightningQubitManager.hpp

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp

pennylane_lightning/core/catalyst/LightningQubitManager.hpp

maliasadi · 2026-01-08T22:39:47Z

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

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

+    // Compact state vector by removing released qubits
+    void CompactStateVector();


minor: this is the style we've been following for a while:

Suggested change

void CompactStateVector();

void compactStateVector();

Should we call this reducedStateVector instead of compact?

Sure! I updated the name here: a5e2132

jzaia18 · 2026-01-08T23:08:25Z

Could you also bump the rc version to 3 within pennylane_lightning/core/_version.py?

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

rniczh · 2026-01-09T01:13:27Z

Could you also bump the rc version to 3 within pennylane_lightning/core/_version.py?

Sure abd373d

jzaia18

Looking good, just 1 last comment

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

dime10

Nice work! Let's get this across the finish line to make sure the dynamic allocation feature actually works :)

Only one question to resolve really:

dime10 · 2026-01-19T20:39:34Z

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

    // maintained,
    //       in particular: measurements must not compute results for released
    //       qubits, only for active/allocated qubits


We can remove the TODO now :)

dime10 · 2026-01-19T20:44:32Z

pennylane_lightning/core/utils/Util.hpp

+ * @param data State vector data to normalize.
+ */
+template <class ComplexT, class Alloc = std::allocator<ComplexT>>
+inline void normalizeStateVector(std::vector<ComplexT, Alloc> &data) {


Just wanted to mention I believe most devices already have a normalize method, since it's needed for the measure implementation:

pennylane-lightning/pennylane_lightning/core/simulators/lightning_qubit/StateVectorLQubit.hpp

Line 908 in 4739a0f

void normalize() {

dime10 · 2026-01-19T21:20:41Z

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

-    // 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.


We should get consensus on the desired behaviour. I don't think this needs to be a big thing, just picking one of the options and checking with product what they think. As far as I know the options are:

(1) it is enforced that qubits must be pure (disentangled) before release

implementation: (simulator-only) raise an error if the user releases an entangled qubit

benefit: nice to catch (likely) programming errors for the user, no undefined or non-deterministic behaviour

issues:

if an entire subset of qubits is released at once, it should be okay if they are entangled with each other, but here we do the release one-by-one so that complicates the detection

takes a bit of work each qubit release

(2) it is well-defined what happens when qubits are not pure:

implementation: reset affected qubits projecting the remainder of state onto of the two subspaces

this could be random (like on hw) or a forced projection onto 0 for example (simulator-only)

since we want the behaviour to be well defined and observable, I think this should happen immediately upon release (e.g. a snapshot would reflect this right away, but I'm not sure that even with option 3 the partial statevector would ever be observable?)

benefit: can be deterministic still, is observable

issues: takes a bit of work each qubit release

(3) it is undefined behaviour what happens when releasing non-pure qubits:

implementation: undefined means we can do whatever is convenient, but at the very least we have to project (and compact) the state before measurement processes and before handing out fresh qubits that reuse previously deallocated space in the statevector

benefit: could be faster to delay than doing something immediately like in the other options

issues: least clean?

Although, even if the projection is done lazily we could still raise an error at that point, it might just be less precise than if we do it immediately. So maybe there is only two choices (error or silent projection), and the question of whether to do it eagerly or lazily is orthogonal to that choice. Another variant could check for an error eagerly, but compact lazily. Or another one could be have a safe release and an unsafe release, which can be toggled by the user.

My gut decision would probably be option 1 for safety (and compact later, and/or safety toggle).

This comment has been minimized.

Sign in to view

dime10 reviewed Jan 6, 2026

View reviewed changes

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp Outdated Show resolved Hide resolved

pennylane_lightning/core/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp Outdated Show resolved Hide resolved

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

rniczh added 2 commits January 8, 2026 13:32

update

54a87e1

fix

14a351e

fix coverage

bf6ff01

rniczh removed the ci:use-gpu-runner Enable usage of GPU runner for this Pull Request label Jan 8, 2026

maliasadi reviewed Jan 8, 2026

View reviewed changes

jzaia18 reviewed Jan 8, 2026

View reviewed changes

pennylane_lightning/core/catalyst/LightningQubitManager.hpp Outdated Show resolved Hide resolved

maliasadi reviewed Jan 8, 2026

View reviewed changes

pennylane_lightning/core/simulators/lightning_gpu/catalyst/LightningGPUSimulator.cpp Outdated Show resolved Hide resolved

maliasadi reviewed Jan 8, 2026

View reviewed changes

rniczh and others added 8 commits January 8, 2026 19:27

Update .github/CHANGELOG.md

203a53b

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

Update .github/CHANGELOG.md

1ed47be

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

remove redundant

81c4859

formatting

6eb8caa

compactStateVector -> reducedStateVector

a5e2132

Update pennylane_lightning/core/catalyst/LightningQubitManager.hpp

43a1b6a

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

formatting

9fc0cc0

bumping to rc3

abd373d

rniczh added 3 commits January 8, 2026 21:31

move normalizeStateVector to Util

60511a8

fix util

128dfff

no need to cover the complex with custom allocator

57b3a61

rniczh mentioned this pull request Jan 9, 2026

Raise error message for dynamic wires with non-prob measurement PennyLaneAI/catalyst#2358

Merged

revert project.toml

0ec104f

rniczh added the ci:use-gpu-runner Enable usage of GPU runner for this Pull Request label Jan 9, 2026

rniczh added 2 commits January 9, 2026 00:27

ci

fcc5229

fix changelog

2b784c7

jzaia18 reviewed Jan 9, 2026

View reviewed changes

pennylane_lightning/core/simulators/lightning_qubit/catalyst/LightningSimulator.cpp Outdated Show resolved Hide resolved

rniczh added 2 commits January 9, 2026 11:35

change fn name of reducedStateVector

0f53498

change to use fn name reduceStateVector for all devices

6c8cdb2

paul0403 requested a review from mlxd January 9, 2026 18:43

jzaia18 removed this from the 0.44.0 milestone Jan 9, 2026

dime10 reviewed Jan 19, 2026

View reviewed changes

		// 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.

Fix dynamic wires with samples #1321

Are you sure you want to change the base?

Fix dynamic wires with samples #1321

Uh oh!

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!

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

maliasadi Jan 9, 2026 • 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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jzaia18 commented Jan 8, 2026

Uh oh!

rniczh commented Jan 9, 2026

Uh oh!

jzaia18 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

dime10 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

rniczh commented Jan 6, 2026 •

edited by sengthai

Loading

codecov bot commented Jan 7, 2026 •

edited

Loading

jzaia18 left a comment •

edited

Loading

maliasadi Jan 9, 2026 •

edited

Loading