PML-388: fixed random phase initialization in entangling layer if trainable is False

[CassNot]

Summary

Non-trainable entangling layers (add_entangling_layer(trainable=False, ...)) were silently pinning every inner/outer phase shifter to 0.0, collapsing the block into a deterministic swap/identity instead of a random fixed unitary. This broke architectures (e.g. reservoir-computing style circuits) that rely on non-trainable entangling layers to inject a fixed random mix. GenericInterferometer now draws random fixed phases for any non-trainable phase shifters at construction time, optionally seeded for reproducibility.

This also fixes a combinatorial bug in NoisyG2SLOSComputeGraph . g2 extra-photon sectors were computed by growing a whole new SLOS graph per input state instead of reusing the base Orthogonal Bad Bits (OBB) partitions, causing incorrect probabilities/tests to rely on the degenerate zero-phase circuits. This bug was highlighted upon solving the first one (the tests were failing because the circuit was not a swap circuit anymore). This fix is inspired by Perceval management of g2 extra photon.

Related Issue

PML-388

Type of change

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

Proposed changes

• merlin/core/components.py: GenericInterferometer gains seed, fixed_inner_values, and fixed_outer_values fields. __post_init__ now draws random phases (via random.Random(seed)) for any non-trainable inner/outer phase shifters instead of hardcoding 0.0.
• merlin/builder/circuit_builder.py: add_entangling_layer(..., seed=None) forwards the seed to GenericInterferometer; the generated Perceval circuit now uses block.fixed_inner_values/fixed_outer_values for non-trainable phase shifters instead of literal 0.0.
• merlin/pcvl_pytorch/noisy_slos.py: Added NoisyG2SLOSComputeGraph._augmented_obb_probs, which reuses the cached OBB partitions of the base input state and grows only the coherent cell's Fock state by the g2-emitted photons, replacing the previous approach of running a full new SLOS graph per extra-photon combination. Fixes incorrect probabilities for g2 extra-photon sectors.
• Test suite updated to reflect that non-trainable layers now carry non-zero random phases, plus new regression tests for reproducibility (seeded vs. unseeded), partial trainable/non-trainable mixes, and the NoisyG2SLOSComputeGraph fix.

How to test / How to run

1. Command lines

pytest tests/builder/test_circuit_builder.py tests/core/test_components.py tests/core/test_merlin_processor_unit.py -v
pytest tests/pcvl_pytorch/test_noisy_g2_slos.py -v

Performance considerations (optional)

No expected performance regression; _augmented_obb_probs reuses cached SLOS subgraphs instead of allocating new ones per combination, so it should be at least as efficient as the previous implementation.

Documentation

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

Checklist

• PR title includes Jira issue key (e.g., PML-126)
• "Related Jira ticket" section includes the Jira issue key (no URL)
• Code formatted (ruff format)
• Lint passes (ruff)
• Static typing passes (mypy) if applicable
• Unit tests added/updated (pytest)
• Tests pass locally (pytest)
• Tests pass on GPU (pytest)
• Test coverage not decreased significantly
• Docs build locally if affected (sphinx)
• With this command: SPHINXOPTS="-W --keep-going -n" make -C docs clean html the docs are built without any warning or errors.
• New public classes/methods/packages are added in the API following the methodology presented in other files.
• Dependencies updated (if needed) and pinned appropriately
• PR description explains what changed and how to validate it

[CassNot]

@claude review

[claude]

Claude finished @CassNot's task in 4m 27s —— View job

---

Claude Code Review — PML-388

I reviewed the diff against origin/release/0.4.1. Overall this is a solid, well-reasoned fix with good docstrings and the random.Random(seed) approach (not touching global RNG state) is a nice touch. The g2 refactor is validated by the existing Perceval-comparison tests in test_noisy_g2_slos.py (test_against_perceval_*, tol 1e-4), which is the strongest possible check. A few things below are worth addressing before merge.

🟠 1. Test-coverage gap for the g2 fix (the riskiest change)

The most complex change (noisy_slos.py, +112 lines, _augmented_obb_probs) has no new regression test in this PR. The diff touches zero files under tests/pcvl_pytorch/ — the PR description's claim of "new regression tests for the NoisyG2SLOSComputeGraph fix" and the pytest tests/pcvl_pytorch/test_noisy_g2_slos.py in How to test don't correspond to anything in the changeset. The fix currently rides entirely on pre-existing Perceval tests.

Recommendation: add a small regression test that pins the augmented extra-photon sector probabilities to the corrected values, so a future refactor can't silently reintroduce the "grow a whole new SLOS graph" behavior. Fix this →

🟠 2. No length validation on user-supplied fixed phase values

merlin/core/components.py:245-252 only populates fixed_*_values when the list is empty. If a caller passes a list shorter than count, circuit_builder.py:954/959/993/998 silently falls back to 0.0 for the missing indices (fixed[i] if i < len(fixed) else 0.0) — reintroducing the exact zero-phase bug this PR fixes, silently. Consider validating that a provided list has exactly count entries (raise ValueError), or at least warn. Fix this →

🟡 3. Lost de-duplication → likely perf regression vs. the claim

_InputStateNoisySLOSComputeGraph.compute_probs (noisy_slos.py:855-866) precomputes probs_per_obb_state once per unique OBB state and reuses it across all partitions/cells. _augmented_obb_probs (noisy_slos.py:246-284) drops this: it calls self._regular_slos_graphs[...].compute_probs(...) fresh for every base_state and every one-hot bad_state in every cell of every partition. The same one-hot state (e.g. [1,0,0]) is recomputed many times. This contradicts the PR's "at least as efficient" note for larger n. A dict memo keyed by the state tuple would restore parity. Fix this →

🟡 4. Behavior change / reproducibility semantics

• Backward compatibility: add_entangling_layer(trainable=False) previously produced a deterministic (near-identity/swap) transform; it now produces a random unitary by default. That's the intended fix, but it silently changes results for any existing user relying on the old behavior — worth an explicit note in the release/changelog (the PR is tagged "Bug fix", not "Breaking change").
• Global-seed reproducibility: because seed=None uses an independent random.Random(), these phases are not reproducible via a global torch.manual_seed(...)/random.seed(...). Users doing reservoir-computing experiments who set a global seed will still get non-deterministic mixes unless they thread an explicit seed= into every add_entangling_layer call. The docstring documents the mechanism, but consider calling out this reproducibility caveat in user-facing docs.

🟢 Minor

• noisy_slos.py:290 output_probs / output_probs.sum(dim=1, keepdim=True) will produce NaN if a sector's total is ever 0. Low risk given probs are positive, but a guard/clamp_min on the denominator would be safer.
• _mzi_factory/_bell_factory capturing block=component as a default arg to avoid closure late-binding is the right pattern — good.
• count = span*(span-1)//2 correctly matches the number of MZI cells in Perceval's RECTANGLE mesh, and the index alignment with fixed[i] in the builder is correct. ✅

Correctness spot-checks (passed)

• The obb cache lookup at noisy_slos.py:408 is safe: slos_graphs_list[0].compute_probs(...) at line 324 always runs first and populates _slos_graph_per_input[tuple(input_state)].
• order == obb.n_photons branch correctly treats the coherent cell as just extra_vec and all original photons as one-hot bad states (the empty base cell is dropped in _generate_obb_partition).
• _regular_slos_graphs index bounds (base_n-1 up to 2·n_photons-1) and _fock_states_per_n keys are all within range.
• to() moves _regular_slos_graphs and guards with hasattr (only defined in the non-g2_distinguishable branch). ✅

---

_{Tasks: read guidelines · reviewed components.py · circuit_builder.py · noisy_slos.py · tests — all complete.}
· branch PML-388-fix-entangling-not-trained

[CassNot]

@claude review

[claude]

Claude finished @CassNot's task in 1m 56s —— View job

---

Claude Code Review — PML-388 (re-review)

Re-reviewed against origin/release/0.4.1 after the latest commits. Since my first pass, two of the four majors are now resolved — nice work:

• ✅ Update issue templates #1 g2 regression test — TestG2AugmentedObbRegression (commit bbf4e72) now pins the extra-photon sector probabilities to Perceval-verified golden values (atol=1e-8), and the docstring documents exactly what the old "grow a new SLOS graph" bug skewed ((1,1,1) → 0.031014 vs. correct 0.024246). This is a strong guard against silent regression. 👍
• ✅ Welcome to merlin Discussions! #2 length validation — GenericInterferometer.__post_init__ (f73657e) now raises ValueError when a user-supplied fixed_inner_values/fixed_outer_values list doesn't have exactly span*(span-1)//2 entries, with parametrized short/overlong/complete-list tests. This closes the silent 0.0 fallback path.

The remaining items are unchanged from last time:

🟡 1. Lost per-state de-duplication in _augmented_obb_probs (perf)

_InputStateNoisySLOSComputeGraph.compute_probs (noisy_slos.py:855-866) precomputes probs_per_obb_state once per unique OBB state and reuses it across all partition cells. _augmented_obb_probs (noisy_slos.py:257, :265) drops this: it calls self._regular_slos_graphs[...].compute_probs(...) fresh for every base_state and every one-hot bad_state in every cell of every partition. The same one-hot state (e.g. [1,0,0]) gets recomputed many times. Functionally correct (the golden test confirms it), but it contradicts the PR's "at least as efficient" note for larger n. A small dict memo keyed by the input-state tuple would restore parity. Fix this →

🟡 2. Reproducibility / backward-compat semantics (docs)

• Backward compatibility: add_entangling_layer(trainable=False) previously produced a deterministic (near-identity/swap) transform and now produces a random fixed unitary by default. That's the intended fix, but it silently changes outputs for anyone relying on the old behavior — worth an explicit line in the release notes / changelog for 0.4.1 (PR is tagged "Bug fix", not "Breaking change").
• Global-seed reproducibility: because seed=None uses an independent random.Random(), these phases are not reproducible via a global torch.manual_seed(...)/random.seed(...). Reservoir-computing users who set only a global seed will still get non-deterministic mixes unless they thread seed= into every add_entangling_layer call. The docstring covers the mechanism; a one-line caveat in user-facing docs would help.

🟢 Minor

• noisy_slos.py:290 output_probs / output_probs.sum(dim=1, keepdim=True) still NaNs if a sector total is ever 0. Low risk (probs are non-negative and sectors are non-empty here), but a clamp_min on the denominator would be a cheap safety net. Note this mirrors the existing :895 normalization, so it's consistent with current behavior — optional.
• key_to_idx.get(key) silently drops any convolved key not in the augmented basis before the renormalization. In practice all convolved keys have exactly n_augmented photons so nothing should be dropped, but as at :892-894 the post-normalization would mask genuine leakage if that assumption ever broke.

Correctness spot-checks (passed)

• random.Random(self.seed) avoids touching global RNG; inner draws first count, outer the next count — deterministic and reproducible under an explicit seed. ✅
• field(default_factory=list) + in-place mutation in __post_init__ — no mutable-default sharing across instances. ✅
• The order == obb.n_photons branch correctly treats the coherent cell as just extra_vec with all originals as one-hot bad states. ✅
• to() moves _regular_slos_graphs under a hasattr guard (only defined in the non-g2_distinguishable branch). ✅
• count = span*(span-1)//2 matches the number of MZI cells in the RECTANGLE mesh, and the builder's fixed[i] index alignment is correct. ✅

Verdict: The two blocking gaps from the first review are addressed and well-tested. The remaining items are a perf parity nit and documentation/changelog notes — none are merge-blocking. Looks good to me. ✅

---

• Read repo guidelines
• Review components.py
• Review circuit_builder.py
• Review noisy_slos.py
• Review test coverage
• Post review
  · branch PML-388-fix-entangling-not-trained

[LF-Vigneux]

Couple of changes in tests, everything else seems sound

[CassNot]

@LF-Vigneux, I do not think it is possible for all values to be 0 because they are samples from a uniform distribution
I added comments in the test files you mentioned

@claude review

[claude]

Claude finished @CassNot's task in 4m 23s —— View job

---

Claude Code Review — PML-388 (3rd pass)

Re-reviewed against origin/release/0.4.1 after f73657e, bbf4e72, 0af228d and the merge.

---

On @LF-Vigneux's "they can technically be 0" — @CassNot is right

Both of you are partly right, and the distinction matters, so here it is precisely:

The reviewer is right that a single value can be exactly 0.0. random.uniform(a, b) is a + (b - a) * random(), and random.random() draws a multiple of 2⁻⁵³ from [0, 1) — it can return exactly 0.0, with probability 2⁻⁵³ ≈ 1.1e-16. So fixed_inner_values[0] == 0.0 is not impossible in the strict sense.

But the assertion as written is any(v != 0.0 for v in ...), not a per-element check. With span=4 there are 4*3//2 = 6 shifters, so the assertion only fails if all six are exactly 0.0:

$$P(\text{fail}) = (2^{-53})^6 \approx 10^{-96}$$

That is not "unlikely," it's beyond any operational notion of possible. Even the stricter all(v != 0.0 ...) form would fail with probability ≈ 6·2⁻⁵³ ≈ 7e-16, i.e. once per ~10¹⁵ runs. Neither is a flake source.

One correction to the phrasing in the added comments (0af228d), though: they say "the probability of getting exactly 0.0 is measure-zero." That's true of the ideal continuous distribution but not of the float64 implementation — 0.0 is genuinely attainable there. The argument that actually holds is the combinatorial one: 6 independent draws all hitting the single value 0.0. Worth rewording so the comment doesn't assert something a careful reader will (correctly) push back on again. Suggested:

# random.uniform() can in principle return exactly 0.0 (probability 2**-53),
# but this only fails if ALL 6 shifters do so simultaneously (~1e-96).
# Guards against the old bug where phases were hardcoded to 0.0.

One more thing in the same spirit, in @CassNot's favour: the range assertion assert 0.0 <= value < 2 * math.pi is also safe, which is less obvious. uniform(0, 2π) computes 2π * r with r ≤ 1 - 2⁻⁵³; the exact product sits 2π·2⁻⁵³ ≈ 6.98e-16 below 2π, and half-ULP at that magnitude is 2⁻⁵¹ ≈ 4.44e-16. Since 6.98e-16 > 4.44e-16 it rounds down, never up to 2π. So the strict < holds unconditionally.

Optional, if you want the argument to be unnecessary rather than merely correct: pass seed= in these three tests. The assertions then become exactly deterministic and the tests still catch the original bug just as well, since a seeded draw is equally non-zero. That would close the discussion by construction rather than by probability. Your call — the current form is not a defect. Fix this →

---

🟠 New: two likely CI lint/format failures

I could not execute ruff here (the sandbox denied the command), so these are from reading the diff against your pyproject.toml config — please confirm locally.

1. tests/core/test_components.py ends without a trailing newline (\ No newline at end of file in the diff, after assert gi.fixed_outer_values == outer). W is in select, so W292 fires.
2. tests/pcvl_pytorch/test_noisy_g2_slos.py:1112 — only one blank line between layer.to("cpu") and class TestG2AugmentedObbRegression:. E is selected and preview = true, which is what enables ruff's E3xx blank-line rules, so E302 fires.

Both are auto-fixed by ruff format. Given the PR checklist has "Code formatted (ruff format)" and "Lint passes (ruff)" ticked, it looks like the last two commits landed after the last format run. Fix this →

🟡 New: _augmented_obb_probs couples to self.device where the rest of the file uses unitary.device

Every other path in noisy_slos.py defensively re-homes onto the unitary's device — probs_regular.to(unitary.device) (:313, :333), probs_one_hot.to(unitary.device) (:325), bit_weight = self._weights[i].to(unitary.device) (:888), convolution.to(unitary.device) (:900), and sector.tensor is allocated on unitary.device (:363).

The new code does not:

• extra_vec is built on self.device (:412) and added to cell[0] (on the OBB's device) at :253;
• output_probs is allocated on self.device (:245) but accumulates conv_probs derived from unitary;
• bit_weight = obb._weights[order] (:247) has no .to(...), unlike its counterpart at :888;
• the caller then does sector.tensor + weight_k * probs where sector.tensor is on unitary.device.

So a unitary on a device other than self.device — which the surrounding code goes out of its way to support — now raises. In the common graph.to("cuda") + cuda unitary flow the two coincide and nothing breaks, which is presumably why tests pass. Cheap fix: use unitary.device throughout _augmented_obb_probs. Fix this →

🟡 New: self._slos_graphs[1:] is now dead in the non-g2_distinguishable branch

__init__ still builds n_photons + 1 NoisySLOSComputeGraphs (:116-127), but after this change compute_probs only ever reads slos_graphs_list[0] — the higher-photon graphs were the old slos_graphs_list[num_photons_added] call site that _augmented_obb_probs replaced. Their constructors are lazy (the OBB helper is built on first compute_probs), so the runtime cost is small, but to() (:470-471) still walks them and they read as live state. Worth collapsing to a single graph so the next reader doesn't assume the extra-photon sectors still route through them.

🟡 Carried over: lost per-state de-duplication (perf)

Unchanged from my last two passes. _InputStateNoisySLOSComputeGraph.compute_probs (:865-876) builds probs_per_obb_state once per unique OBB state and reuses it across all partition cells. _augmented_obb_probs (:257, :265) calls compute_probs fresh for every base_state and every one-hot bad_state in every cell of every partition — the same one-hot state (e.g. [1,0,0]) is recomputed many times. Functionally correct (the golden test confirms it), but it contradicts the PR's "at least as efficient" note as n grows. A dict memo keyed by the state tuple restores parity. Fix this →

🟢 Minor

• Seed semantics shift with trainability. In components.py:270-279 a single rng feeds inner then outer. With trainable_inner=True, the outer phases consume draws 0..count-1 instead of count..2count-1, so seed=42 yields different outer phases depending on trainable_inner (and on whether fixed_inner_values was supplied). Harmless in isolation, but surprising if someone compares a fully-fixed layer against a half-trainable one under "the same seed." Drawing both blocks unconditionally, or deriving two RNGs, would decouple them.
• fixed_*_values + trainable_*=True is length-validated but then silently ignored by the builder (which emits pcvl_module.P(...) for the trainable side). A warning would be kinder than silence.
• Unseeded instances now compare unequal. GenericInterferometer is an eq=True dataclass, so two identically-configured unseeded blocks are no longer ==. I checked MerlinModule — it uid's via uuid4() (module.py:91), not component equality — so nothing in-tree breaks. Noting it only in case downstream code dedupes on component equality.
• noisy_slos.py:290 — output_probs / output_probs.sum(dim=1, keepdim=True) still NaNs on a zero sector total. Mirrors the existing :905 normalization, so it's consistent; a clamp_min would be a cheap net.
• key_to_idx.get(key) (:283) silently drops convolved keys outside the augmented basis before renormalization. All convolved keys should have exactly n_augmented photons so nothing drops in practice, but as at :902-904 the post-normalization would mask genuine leakage if that ever broke.
• Changelog: add_entangling_layer(trainable=False) goes from deterministic to random-by-default. Intended, but it silently changes outputs for existing users — worth an explicit line in the 0.4.1 release notes even though the PR is tagged "Bug fix". Also worth a user-docs note that seed=None uses an independent random.Random() and is therefore not reproducible via a global torch.manual_seed/random.seed.

✅ Verified this pass

• count = span*(span-1)//2 matches the MZI cell count of Perceval's RECTANGLE mesh; the builder's fixed[i] indexing aligns. The span > 1 guard makes count = 0 for degenerate spans, and the validation message then correctly reads "exactly 0 entries".
• field(default_factory=list) + in-place assignment — no mutable-default sharing.
• block: GenericInterferometer = component as a default arg to defeat closure late-binding is the right pattern in both _mzi_factory and _bell_factory.
• _regular_slos_graphs indexing is sound: built for n_i in range(1, 2*n_photons+1), so [base_n - 1] → n = base_n and [0] → n = 1. ✅
• self._fock_states_per_n spans 1..2*n_photons (:131-136), covering every n_augmented and base_n reachable here. ✅
• obb = slos_graphs_list[0]._slos_graph_per_input[tuple(input_state)] (:419) is safe: input_state is typed list[int] | tuple[int, ...] (:297), so the tuple key is hashable-by-value, and the :328 call always populates the cache first via :662-675.
• order == obb.n_photons correctly treats the coherent cell as just extra_vec with all originals as one-hot bad states.
• to() guards _regular_slos_graphs with hasattr (only defined in the non-g2_distinguishable branch). ✅
• TestG2AugmentedObbRegression is a genuinely strong guard — golden values at atol=1e-8, with the pre-fix skew documented inline ((1,1,1) → 0.031014 vs. correct 0.024246).

---

Verdict: The core fix is sound and the two blockers from my first pass stay closed. The reviewer's concern does not hold — the any() guard makes those assertions deterministic in practice — though the "measure-zero" wording in the new comments overstates the case and is worth a small reword. Before merge I'd fix the two ruff issues (likely red CI) and the unitary.device coupling; the rest are non-blocking.

I was not able to run pytest or ruff in this environment (both commands require approval), so the lint findings are from reading the diff and the test results are unverified here.

• Read repo guidelines (AGENTS.md, pyproject.toml)
• Re-review diff vs origin/release/0.4.1
• Assess the "can technically be 0" question
• Check status of prior review items
• Post review
  · branch PML-388-fix-entangling-not-trained

[LF-Vigneux]

I approve