PML-380: Unitary device fix for g2

[LF-Vigneux]

Summary

Fixed this issue:

Noisy g2 path allocates on self.device/self.dtype, not unitary.device — merlin/pcvl_pytorch/noisy_slos.py:238-258,316. p_emit/weight_k and each SectorResult.tensor are built on self.device, while the sub-graph probs live on unitary.device. A CUDA unitary on a graph constructed with default device=None (without a prior .to()) raises a cross-device error. The non-g2 path correctly derives device from unitary.device (:768) — mirror that here.

Related Issue

PML-380

Type of change

• Bug fix
• New feature
• Documentation update
• Refactor / Cleanup
• Performance improvement
• CI / Build / Tooling
• Breaking change (requires migration notes)

Proposed changes

• Initialized the g2 and sectors results with the unitary's device instead of the SLOS'.

How to test / How to run

pytest -q

Documentation

• User docs updated (Sphinx)
• Examples / notebooks updated
• Docstrings updated
• Updated the API

[LF-Vigneux]

@claude review

[claude]

Claude finished @LF-Vigneux's task in 1m 8s —— View job

---

Code Review — PML-380: Unitary device fix for g2

• Read the diff and surrounding code
• Verify the fix against the non-g2 reference path
• Check for remaining cross-device tensor allocations
• Review the added test
• Post review

---

Summary

The fix is correct, minimal, and well-targeted. The two changed allocations in NoisyG2SLOSComputeGraph.compute_probs now derive their device from unitary.device instead of self.device, matching the pattern the non-g2 path already uses. This resolves the cross-device error that occurred when a CUDA unitary is passed to a graph built with device=None.

What I verified

• merlin/pcvl_pytorch/noisy_slos.py:238 — _g2 now on unitary.device, so p_emit and each weight_k (line 242) live on the unitary's device. ✅
• merlin/pcvl_pytorch/noisy_slos.py:254 — each SectorResult.tensor is zeroed on unitary.device. ✅
• The downstream combinations are now consistent:
  • sector.tensor + weight_k * probs_regular (line 261) — both operands on unitary.device (probs_regular comes from the subgraph, which already runs on the unitary's device).
  • sector.tensor + weight_k * probs (line 316) — probs comes from convolve_distributions/subgraph outputs, all on unitary.device.
• No remaining self.device allocations in compute_probs that would feed into the sector tensors. The one in _get_extra_photon_combinations (noisy_slos.py:173) is only used for .tolist() enumeration, so its device is irrelevant.
• SectorResult.__post_init__ builds keys as plain Python tuples (no device), so nothing else needs moving.

Minor observations (non-blocking)

1. Test is CUDA-gated, so it won't run in CPU CI. test_noisy_g2_slos_unitary_cuda_graph_cpu is skipped when CUDA is unavailable (tests/pcvl_pytorch/test_noisy_slos.py:369). This is unavoidable — the bug is inherently a cross-device mismatch that can't manifest on CPU-only — but it does mean the fix is not exercised by a typical CPU CI run. Worth being aware of.

2. Assertion could be stronger. The test only checks probs.sectors[0].tensor.device (test_noisy_slos.py:380). Since g2=0.7 with n_photons=3 produces multiple sectors, asserting every sector's device would give slightly broader coverage:

   assert all(s.tensor.device == unitary.device for s in probs.sectors)

3. g2=0.7 is outside the documented valid range. The docstring at noisy_slos.py:236 states the p = ((1 - g2) - sqrt(1 - 2*g2)) / g2 conversion is valid for g2 ∈ [0, 0.5]. With g2=0.7, _disc clamps to 0, so the code still runs and returns a device-correct tensor — fine for a device-focused test — but a value within [0, 0.5] would be a more representative fixture.

Verdict

Looks good to merge. The core change is sound and mirrors the existing correct behavior. The observations above are optional test-quality nits, not correctness issues.
• PML-380-unitary-g2-device

[CassNot]

The fix does not work (cuda issue) once the perceval typo and type issue are solved (n=5 and dtype=np.complex64)

You will find my fix in the file attached:

• lines 214-215 and 225-226 and 232-238: I am ensuring the probs stay on the same device as the unitary (if g2_distinguishable or not)
• lines 783 and 796, I am also sending the bit_weight and convolution to the unitary.device

So now, all probability tensors are explicitly placed on the unitary's device before use. This should prevent device mismatch errors during convolution and other operations.

With all these fixes, it works on GPU

noisy_slos.py

[CassNot]

here it should be n=5 and not m=5 otherwise it fails with a perceval error

[CassNot]

Another error from this line is that the type is np.complex128 which is inconsistent with the torch.float32 defaults types from the compute graph. If you set it, you need to change the type of the graph. My fix would be to just have the dtype=np.complex64 here and that works !

[CassNot]

@claude I believe your analysis was wrong and the code was not working on GPU because

• when compute_probs is called with a unitary on a different device, there's a device mismatch
• when computing one-hot state probabilities, they're not being moved to the unitary's device before convolution
  All probability tensors should be explicitly placed on the unitary's device before use, preventing device mismatch errors during convolution

[claude]

Claude finished @CassNot's task in 2m 12s —— View job

---

Re-analysis — you're right, my earlier review was incomplete

• Re-read NoisyG2SLOSComputeGraph.compute_probs and the subgraph OBB path
• Trace every tensor that flows into the convolution / accumulation
• Re-check the test against Perceval's API and the graph dtype
• Post corrected conclusion

---

You're correct, and my original review missed the real bug. I verified only the two allocations inside NoisyG2SLOSComputeGraph.compute_probs (_g2 and each SectorResult.tensor) and concluded the fix was complete. I did not trace into the subgraph that actually produces probs_regular / the one-hot probs — and that's where the mismatch lives.

Where it actually breaks

NoisySLOSComputeGraph.compute_probs (the OBB / partial-indistinguishability path, noisy_slos.py:771-785):

for i, partition in enumerate(self._partitions):
    bit_weight = self._weights[i]          # built on self.device at construction (:701-704)
    ...
    _, convolution = convolve_distributions(fock_states, *cell_distributions)  # on unitary.device
    output_probs += bit_weight * convolution * count.item()   # CPU * CUDA  -> device mismatch

self._weights (and self.g/self.b) are materialized at construction time on self.device (noisy_slos.py:697-704). With device=None and no prior .to() — exactly the test scenario — they stay on CPU while convolution/output_probs are on the unitary's CUDA device. So bit_weight * convolution at :785 raises the cross-device error.

Because the g2 graph gets probs_regular and the one-hot distributions from this subgraph (:211-223), the failure happens before my "fixed" outer allocations are ever combined. My PR-fix change (_g2 and sector tensors → unitary.device) is necessary but not sufficient — your additions moving bit_weight and the convolution onto unitary.device are the part that actually resolves it. Your diagnosis (#1 device mismatch when the unitary is on another device, #2 one-hot probs not on the unitary's device before convolution) is accurate.

The test never actually exercised the fix

Both of your review comments on test_noisy_slos.py:375 are correct, and together they mean the test raised before reaching the device assertion — so it never validated anything, which is why my "looks good" was wrong to rely on it:

1. pcvl.Matrix.random_unitary(m=5) — random_unitary's dimension argument is n, not m; m=5 is an unexpected keyword → Perceval error. Should be random_unitary(5) / n=5.
2. dtype=np.complex128 — the graph defaults to torch.float32 (→ complex64), so a complex128 unitary is a dtype mismatch. dtype=np.complex64 is the right fixture (or set the graph dtype to match).

Corrected verdict

Changes requested is right. The complete fix needs:

• the two outer allocations in NoisyG2SLOSComputeGraph.compute_probs (already in this PR), plus
• bit_weight and the convolution result placed on unitary.device in NoisySLOSComputeGraph.compute_probs (your noisy_slos.py attachment), and
• the test fixture corrected to n=5 + np.complex64 so it actually runs on GPU.

Thanks for catching this — apologies for the incomplete first pass. Want me to apply your attached noisy_slos.py changes plus the test fixes to this branch, or will you push them yourself?
• PML-380-unitary-g2-device