PML-376: hints on FF with PartialMeasurement

[CassNot]

Summary

I am adding an explanation in the documentation of the FeedForward in MerLin to reflect the findings of #274
The goal is to explain to users how to use PartialMeasurement when what they want to do with the FeedForwardBlock is not possible (here, the input parameters cannot be used inside a branch)

Related Issue

PML-376

Type of change

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

Proposed changes

Summary of what was added, right after the main FeedForwardBlock example and before "Further Reading":

• A disclaimer note pointing to issue Feedforward block cannot detect parameters outside initial layer. #274 and stating this is fixed in MerLin 0.5 (closing Feedforward block cannot detect parameters outside initial layer. #274 ).
• An explanation of the exact limitation (input parameters must be consumed by the first pre-measurement stage) and the ValueError users hit, with a minimal repro.
• A workaround showing how to reproduce branch-local classical inputs manually with MeasurementStrategy.partial(...), routing each PartialMeasurement branch into its own QuantumLayer.

How to test / How to run

1. Command lines

make -C docs html

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 5m 16s —— View job

---

Review: PML-376 — hints on FF with PartialMeasurement

• Gather context (diff, AGENTS.md, related source)
• Verify the documented limitation & API against the source
• Review prose and code snippets
• Post review

Docs-only change, well placed (right after the main example, before Further Reading) and the diagnosis is correct. I verified the claims against the source:

• The quoted ValueError is real — merlin/algorithms/feed_forward.py:695-700, covered by tests/algorithms/test_feed_forward_block.py:598-607, and branch-circuit parameters are discovered by experiment.get_circuit_parameters(), so the failure really is "not used in the first stage" and not "no matching prefix".
• MeasurementStrategy.partial(modes=..., computation_space=...) (merlin/measurement/strategies.py:400), PartialMeasurement.branches / branch.outcome / .probability / .amplitudes (merlin/core/partial_measurement.py:37-115), ComputationSpace.FOCK, and QuantumLayer.output_keys (merlin/algorithms/layer.py:1592) all check out. return_object=True is correctly not needed here: PARTIAL always returns a PartialMeasurement (merlin/algorithms/layer.py:1293-1301).
• The :doc: target /quantum_expert_area/partial_measurement resolves.

Main concerns are in the workaround snippet, which is the part users will copy.

---

1. The workaround as written cannot be trained (highest impact)

docs/source/user_guide/feed_forward.rst — the branch QuantumLayers are constructed inside the per-branch loop, i.e. on every forward pass. That means trainable_parameters=["A"]/["B"] are freshly re-initialized each call, so nothing learns and no optimizer can hold them. Since the whole point of this workaround is training a feedforward model, the example should show the training-safe shape:

• build partial_layer and the branch layers once in an nn.Module.__init__ (e.g. an nn.ModuleDict keyed by outcome), and
• in forward, do branch_layer.set_input_state(branch.amplitudes) then branch_layer(x).

set_input_state accepts a StateVector (merlin/algorithms/layer.py:902-947), and it's the only option here because forward() explicitly refuses to mix a StateVector with float tensors: "Cannot mix torch.Tensor and StateVector inputs in the same forward() call. … set it via the constructor or set_input_state()" (merlin/algorithms/layer.py:1113-1118).

This is also what FeedForwardBlock itself does — stage layers are built at construction (_build_stage_runtime) and amplitudes are passed at run time (merlin/algorithms/feed_forward.py:1023-1037). Fix this →

2. Batch handling is silently wrong for batch_size > 1

branch_output = branch_layer(x).squeeze(0) if input_parameters else branch_layer().squeeze(0)
...
probabilities[output_key] = branch.probability.squeeze(0) * branch_output[index]

With batch > 1, branch_layer(x) is (batch, n_keys), .squeeze(0) is a no-op, and branch_output[index] then indexes the batch dimension — plausible-looking but wrong numbers, no error. Prefer the batch-safe form (branch_output[:, index], and broadcast branch.probability, e.g. branch.probability.unsqueeze(-1) / index it), or state explicitly that the snippet is written for batch_size == 1. Worth also saying how batch flows: the prefix layer is called with no features (batch 1) while the branch layers get x.shape[0], and the workaround relies on broadcasting between the two. Fix this →

3. Key reconstruction assumes the measured mode is mode 0

output_key = (branch.outcome[0], *remaining_key) only produces correct full-system Fock keys because exactly one mode is measured and it happens to be the first one. Either say so, or build the key from partial_measurement.measured_modes / unmeasured_modes so readers generalizing to modes=[1, 2] don't get scrambled keys.

4. The ValueError repro is not runnable

experiment is never created, there's no prefix circuit, no pcvl.Detector.pnr(), and no input_state= (required — merlin/algorithms/feed_forward.py:678-679). Copy-pasting gives NameError; adding a bare experiment without a detector gives a different error ("requires detectors preceding the first feed-forward provider"), which undercuts the point. tests/algorithms/test_feed_forward_block.py:560-607 is a complete, already-verified repro that could be trimmed down here. Fix this →

5. "exactly what FeedForwardBlock does internally" is overstated

Internally the block builds its layers once, chains stages via MeasurementStrategy.amplitudes(), and applies its own per-stage detector transform (_build_partial_detector). Suggest "closely mirrors" instead of "is exactly", plus one sentence that this workaround covers a single feedforward stage — multi-stage experiments require nesting the same pattern.

6. Snippet is unverified

Nothing in CI executes RST code blocks (no doctest/codeblocks runner in .github/), so these snippets can rot silently. I could not execute Python in this sandbox (permission denied), so I did not run the example end-to-end. Please confirm it runs — in particular the (2,) branch, which feeds a zero-photon conditional state into QuantumLayer(circuit=Circuit(2), ...) — and that the six reconstructed keys sum to 1.

Minor / nits

• Source typo, quoted verbatim: merlin/algorithms/feed_forward.py:699 reads "Create you own stages" (missing "r", no final period). Worth fixing at the source and updating the quote. Fix this →
• Import paths: partial_measurement.rst uses the public from merlin import QuantumLayer; this snippet uses merlin.algorithms.layer. Not wrong (the existing example in this file also uses submodule paths), but the public path is friendlier.
• pcvl.Unitary.random(3) makes the example non-reproducible; a fixed pcvl.BS()-based prefix would be better if you ever add expected output.
• Precision: the limitation is about prefixes (startswith matching, merlin/algorithms/feed_forward.py:624-644), and a related failure mode is params covered by no spec at all ("Parameters ... not covered by any trainable or input spec"). One clause on prefix matching would pre-empt confusion.
• Version promise: "expected to be solved in MerLin 0.5" will rot — consider a # TODO-style tracking note or Jira reminder to delete this section when the new backend lands.
• Style: this heading is Title Case while the rest of the file is sentence case ("FeedForwardBlock in MerLin", "Further Reading"); the underline is also longer than the title (valid RST, just inconsistent).

No blocking issues in the explanation — the limitation is described accurately. Items 1 and 2 are what I'd want fixed before merging, since a docs snippet that can't train and silently mis-indexes batches will generate support questions.

I did not change any files (review only).
• branch PML-376-doc-ff-pm

[CassNot]

I agree with Claude's comment that my fixed proposition should have been verified.
I created tests and updated the explanation to include a 3-mode circuit which clarified the explanations.
The fix is verified and ready for a 2nd @claude review

[claude]

Claude finished @CassNot's task in 5m 7s —— View job

---

Review (round 2): PML-376 — hints on FF with PartialMeasurement

• Gather context (new diff, new test file, AGENTS.md)
• Verify the updated workaround against the source
• Check whether the new tests actually pin the documented snippet
• Re-check the 6 items from round 1
• Post review

Thanks for adding tests — the 3-mode rewrite does make the explanation clearer, and the training-safe shape (build once in __init__, set_input_state at runtime) and the batch-safe indexing (branch_output[:, index]) are both correct now. Round-1 items #1, #2 and #5 are addressed.

⚠️ I could not execute anything in this sandbox — python, pytest, ruff and sphinx all came back "This command requires approval". Everything below is static reading of the source. Docs Validation is green; CI Tests was still in progress when I looked (run).

The headline problem: the snippet in the .rst is not the code that the test file verifies. They diverged, and the doc version is the broken one — while the new note tells readers it is verified.

---

1. The documented snippet references an undefined name (blocking)

docs/source/user_guide/feed_forward.rst, inside BranchFeedforward.forward:

if sum(branch.outcome) == sum(input_state):   # <-- global, not self

input_state is not a parameter of forward and is never stored on self. It only resolves because the Usage block below happens to define a module-level input_state = [1, 1, 0]. A reader who copies just the class, or names their state anything else, gets NameError — or worse, silently picks up an unrelated global. The test got this right (sum(self.partial_layer.input_state), line 96); port that into the doc.

2. The documented snippet's own multi-batch assertion crashes (blocking)

Doc, vacuum branch:

probabilities[(branch.outcome[0], 0, 0)] = branch.probability   # shape (1,)

The prefix layer runs with no features, so branch.probability has batch 1, while every other entry is (batch_size,). The doc's own verification block then does:

for i in range(x.shape[0]):                     # 3
    batch_total = sum(p[i].item() for p in probs_multi.values())

→ IndexError at i = 1 on the (2, 0, 0) entry. The test file added .expand(x.shape[0] if x is not None else 1) (lines 97-99) precisely to fix this; the doc never got it. So the printed example cannot run as printed for batch_size = 3. Fix this →

3. Doc re-declares the branch config inside the loop

branch_config = { (0,): ["x"], (1,): ["x"] }   # rebuilt every branch, every call
input_params = branch_config[branch.outcome]

Duplicates branch_configs from __init__, and has no (2,) entry — any non-vacuum branch without input parameters would KeyError. The test's self.branch_configs_dict[branch.outcome] is the right shape.

Suggestion for 1-3: use .. literalinclude:: ../../../tests/algorithms/test_feedforward_manual_workaround.py with :pyobject: BranchFeedforward. That makes "verified by the test suite" structurally true instead of a claim that already went stale within one commit.

4. The new error test passes for the wrong reason (blocking)

tests/algorithms/test_feedforward_manual_workaround.py:310-316 passes the same prefixes in both specs:

trainable_parameters=["A_branch", "B_branch"],
input_parameters=["A_branch", "B_branch"],

_map_prefix_to_params assigns both names in the trainable loop, then the input loop hits them again and raises at merlin/algorithms/feed_forward.py:633-637:

ValueError: Parameter 'A_branch' matches multiple prefixes: 'A_branch' and 'A_branch'

That fires long before _build_stage_runtime ever reaches the documented check at feed_forward.py:695-700. Because pytest.raises(ValueError) has no match=, the test is green while asserting nothing about the limitation it claims to cover — and the documented error is never exercised. Note the doc repro is right here (trainable_parameters=["A","B"], input_parameters=["x"] with a shared x), so the test just needs to mirror it, plus match="The first stage must use all of the input parameters". Also worth noting tests/algorithms/test_feed_forward_block.py:598-607 already covers this exact error with the match=, so this test is either redundant or should be replaced by a check of the doc's own 4-mode repro. Fix this →

5. "Verified" is weaker than the doc claims — nothing compares against FeedForwardBlock

The doc asserts:

This manual composition reproduces the same probabilities as FeedForwardBlock (for single-stage experiments)

No test checks that. The three workaround tests assert dict/shape/range, sum-to-1, and that parameters move. Sum-to-1 is satisfied by any per-branch normalized distribution weighted by branch probabilities — it does not distinguish a correct chained circuit from a wrong one, and it wouldn't catch the key-scrambling in #6.

An equivalence test would: build an experiment FeedForwardBlock does accept (input parameter consumed by the prefix stage, or trainable-only branch params), run both, torch.testing.assert_close on the matching keys. Until that exists, I'd soften the wording to "reproduces the same structure". Two real differences to mention while you're there: FeedForwardBlock chains stages with MeasurementStrategy.amplitudes() and applies its own per-stage detector transform (_build_partial_detector, feed_forward.py:539-551), whereas the workaround uses MeasurementStrategy.partial(..., FOCK) with no detector model — the two coincide only for PNR detectors.

6. test_manual_feedforward_keys_arbitrary_measured_modes doesn't test what its name says

The name and docstring (line 320-321) promise generalization beyond mode 0, then line 321 concedes "For demonstration, we use measured_modes=[0]". The assertion key[0] in (0, 1, 2) is trivially true because output_key = (branch.outcome[0], *remaining_key) puts branch.outcome[0] first by construction — it re-asserts the code, not the semantics.

Worse, the comment at line 122 is wrong:

# This works for arbitrary measured_modes because we use the branch outcome directly.
output_key = (branch.outcome[0], *remaining_key)

With measured_modes=[1] the correct full-system key is (remaining[0], outcome[0], remaining[1]); this emits (outcome[0], remaining[0], remaining[1]) — scrambled, silently. Either build the key from partial_measurement.measured_modes / unmeasured_modes and make the test actually use measured_modes=[1] (it would then catch the bug), or delete this test and keep the honest single-measured-mode caveat that's already in the doc. Per AGENTS.md ("do not weaken tests", "do not broaden assertions"), a test that asserts less than its name claims is worse than no test. This is round-1 item #3, still open.

7. Unseeded randomness → possible flake

All five tests use pcvl.Unitary.random(3) / random(4) with no seed, and tests/conftest.py sets no global seed. The existing test_feed_forward_block.py deliberately avoids Unitary.random entirely (0 occurrences) in favour of named-parameter interferometers. Sum-to-1 is basis-independent so those two are safe, but test_manual_feedforward_workaround_trainable asserts every parameter moved by more than atol=1e-6 after 5 SGD steps at lr=0.01 — a draw where one branch has near-zero probability, or where the single retained key per branch is insensitive to that branch's A/B, will make it fail. Add torch.manual_seed(...) + a seeded/fixed prefix (e.g. pcvl.BS()-based) so failures are reproducible.

Related: the dict comprehension at lines 261-265 keeps only the last key per key[0], so the loss depends on dict ordering. Intentional? A sum over all keys of the branch would be more robust.

---

Minor / nits

• Doc repro vs. workaround are different systems — the ValueError example is 4 modes / Circuit(3) branches, the workaround is 3 modes / Circuit(2) branches. Adding "the workaround below uses a smaller 3-mode system for readability" would avoid readers hunting for the mapping.
• Dead code in both copies: branch_prob_weighted = branch.probability.unsqueeze(-1) immediately followed by branch_prob_weighted.squeeze(-1) * branch_probs_for_key is a no-op round trip — just branch.probability * branch_probs_for_key. The doc's own bullet ("use unsqueeze(-1) or indexing") is therefore documenting code that does nothing.
• forward(self, x=None) with x=None and a branch that needs input parameters calls branch_layer(None) → TypeError: Unsupported input types: ['NoneType'] (layer.py:1058-1068). Either drop the default or guard it.
• Formatting: ~20 lines carry trailing whitespace (24, 26, 59, 141, 145, …) and the imports aren't isort-ordered (torch before perceval), so ruff check tests/ would report W293 / I001. CI only lints merlin/ (.github/workflows/ci.yml:56-57), so this won't fail the build, but the PR checklist has "Code formatted (ruff format)" unchecked — worth running.
• circuit, _trainable_params, input_params = self.branch_configs_dict[...] (line 109) — circuit is unused; _circuit for consistency with _trainable_params.
• Round-1 nits still open, all cosmetic: source typo "Create you own stages" (feed_forward.py:699) still quoted verbatim in the doc; merlin.algorithms.layer instead of the public from merlin import QuantumLayer; heading in Title Case while the rest of the file is sentence case, with an over-long underline; and the "expected to be solved in MerLin 0.5" line will rot (a # TODO/Jira reminder to delete this section would help).
• Doc comment "n_photons (quantum dimension)" — it's a photon count, not a dimension.

---

Summary. The prose and the diagnosis are good, and the workaround's shape is now right. But the two blockers are that the snippet users will copy still can't run (items 1-2, undefined global + batch-shape mismatch in the doc's own assertion), and the new error test doesn't reach the error it claims to (item 4). Item 6 is a correctness comment that is actively false. Item 5 is about the strength of the word "verified": the tests show the pattern runs and normalizes, not that it matches FeedForwardBlock.

I did not change any files (review only), and I was unable to run the test suite, ruff, or the docs build in this sandbox — please rely on the CI run for green/red.
• branch PML-376-doc-ff-pm

[CassNot]

Indeed, good comments
@claude review

[claude]

Claude finished @CassNot's task in 5m 51s —— View job

---

Review (round 3): PML-376 — hints on FF with PartialMeasurement

• Gather context (latest diff, AGENTS.md, CI status)
• Re-check round-2 blockers (Update issue templates #1 undefined global, Welcome to merlin Discussions! #2 batch shape, Docs #3 config duplication, include license in py files #4 wrong-reason error test, code formatting #6 measured-modes)
• Verify the key-reconstruction fix against the source
• Check doc snippet ↔ test file parity
• Post review

Good progress — most of round 2 is genuinely fixed, and one of the fixes is better than I expected. Verified statically:

• Round-2 Update issue templates #1 fixed — sum(self.partial_layer.input_state) replaces the module-level global in both copies.
• Round-2 Welcome to merlin Discussions! #2 fixed — the vacuum branch now .expand(x.shape[0] if x is not None else 1), so all entries share the batch dim and the doc's own for i in range(x.shape[0]) assertion can run.
• Round-2 Docs #3 fixed — forward() reads self.branch_configs / self.branch_configs_dict built in __init__ instead of rebuilding a dict per branch.
• Round-2 include license in py files #4 fixed — trainable_parameters=["A","B"] / input_parameters=["x"] no longer collide in _map_prefix_to_params (feed_forward.py:600-645), the prefix Unitary.random(4) exposes no named params so input_params_set is empty ≠ {"x"}, and the match= now pins feed_forward.py:695-700. The test reaches the error it claims.
• Round-2 move tests to root, add license header in all files #5 fixed — wording softened to "same output structure", with the PNR/multi-stage caveat spelled out. Accurate now.
• Round-2 code formatting #6 — the code fix is correct and more general than the doc claims. I traced the ordering: measured_modes / unmeasured_modes are both built in ascending mode index (merlin/measurement/detectors.py:184-189), and branch.outcome / remaining_key are extracted in exactly that same order (detectors.py:476-477). So zip(measured_modes, branch.outcome) + zip(unmeasured_modes, remaining_key) reconstructs correct full-system keys for any measured-mode set, even if the user passes modes=[2, 0]. The closing note's "ensure key reconstruction uses the correct measured modes" undersells this — the snippet already handles it.

⚠️ I could not execute anything in this sandbox — python, pytest, ruff all returned "This command requires approval", and find outside the repo is blocked. Everything below is source reading. Docs Validation is green; CI Tests (run 30575124918) was still in progress.

---

1. The ValueError repro in the doc almost certainly raises the wrong error (main remaining issue)

docs/source/user_guide/feed_forward.rst:163-165:

provider = pcvl.FFCircuitProvider(1, 0, Circuit(3))          # default branch: 3 modes
provider.add_configuration([0], pcvl.BS(pcvl.P("x")) // pcvl.BS(pcvl.P("A")))   # 2 modes
provider.add_configuration([1], pcvl.BS(pcvl.P("x")) // pcvl.BS(pcvl.P("B")))   # 2 modes

BS // BS is a 2-mode circuit, but the provider's default is Circuit(3). Every other FFCircuitProvider in this repo keeps the configuration and default circuits at the same width — test_ff_perceval_integration.py:26-34, test_feed_forward_block.py:75-99 / :482-484 / :588-590, the existing example at feed_forward.rst:99-100, and your own new test (Circuit(m - 1) for both, line 313). A provider whose branches have a different footprint than its default has no consistent placement, so this should be rejected by Perceval before FeedForwardBlock is ever reached.

I could not run it to confirm which error surfaces (no Python, and Perceval's source is outside the readable tree) — but this is the one code block in the PR that the new test file does not cover, and it's the block whose entire purpose is to produce a specific ValueError. If it raises anything else, the section teaches the wrong diagnosis. Please paste the actual output of running it.

Related, smaller: the doc creates two distinct pcvl.P("x") objects (one per configuration) while the test shares a single instance (test_feedforward_manual_workaround.py:306). Whether same-named distinct Parameter objects are unified is a Perceval implementation detail; the test's shared-instance form is the one you've actually exercised, so mirror it. Fix this →

2. The doc snippet and the test are still not the same code, but the doc says it is verified

feed_forward.rst closes with:

The pattern is verified by the test suite in tests/algorithms/test_feedforward_manual_workaround.py, which covers single-batch, multi-batch, trainability, and error cases.

The two BranchFeedforward classes differ: the doc's takes (input_state, prefix_circuit) and hardcodes modes=[0] plus the branch circuits inside __init__; the test's takes (input_state, prefix_circuit, branch_configs, measured_modes) and factors the reconstruction into a reconstruct_full_key helper. The logic matches this commit — but that was also true one commit ago, right before they drifted. And the "error cases" clause is not accurate: the doc's error snippet (item 1) is a different experiment from the test's.

.. literalinclude:: ../../../tests/algorithms/test_feedforward_manual_workaround.py with :pyobject: BranchFeedforward would make the claim structurally true instead of a promise that has already broken once. Fix this →

3. test_manual_feedforward_keys_arbitrary_measured_modes still cannot fail

tests/algorithms/test_feedforward_manual_workaround.py:350-357 — good that it now passes measured_modes=[1], that's the right move. But the assertion is:

for key in probabilities:
    assert key[1] in (0, 1, 2)

The old, scrambled implementation (output_key = (branch.outcome[0], *remaining_key)) would put remaining_key[0] at index 1 — also always in (0, 1, 2) for a 2-photon system. So this test passes for both the correct and the buggy code. It's guarding nothing.

A one-line discriminating assertion exists here. With input_state=[1,1,0] and measured_modes=[1], the vacuum branch (outcome (2,)) must land at key (0, 2, 0); the scrambled version emits (2, 0, 0):

assert (0, 2, 0) in probabilities
assert (2, 0, 0) not in probabilities

Stronger still, and cheap: with identity branch circuits, the manual composition must reproduce a plain full measurement of the prefix — build QuantumLayer(circuit=prefix, input_state=input_state, measurement_strategy=MeasurementStrategy.probs(ComputationSpace.FOCK)) and torch.testing.assert_close key-by-key. That single test pins the key mapping and the probability weighting, and would have caught the round-2 scrambling. Per AGENTS.md ("do not broaden assertions"), the current form is worse than no test because its name advertises coverage it doesn't provide. Fix this →

4. Coverage gap: nothing tests gradient flow through set_input_state

test_manual_feedforward_workaround_trainable is the test that justifies the doc's central claim ("this is the only safe way to preserve trainable parameters"). But the prefix in every test is pcvl.Unitary.random(3) — a fixed unitary with no trainable parameters — so model.named_parameters() contains only the branch A and B. The riskiest link in the pattern is whether gradients propagate back from a branch layer, through set_input_state(branch.amplitudes) → _embed_amplitude_tensor → computation_process.input_state (layer.py:943-948, :881-888), into a prefix trainable parameter. That is exactly what the test does not exercise.

Adding a named parameter to the prefix circuit (e.g. one pcvl.PS(pcvl.P("theta")), with trainable_parameters=["theta"] on partial_layer) and asserting theta.grad is not None and theta.grad.abs().sum() > 0 would close this. It's also the assertion most likely to catch a real regression.

5. Unseeded randomness — flake risk (round-2 #7, still open)

All five tests use pcvl.Unitary.random(3) / random(4) with no seed, and tests/conftest.py sets none. test_feed_forward_block.py deliberately uses zero Unitary.random in favour of named interferometers. Sum-to-1 and structural assertions are basis-independent, so those are safe — but test_manual_feedforward_workaround_trainable:282-286 asserts every parameter moved by more than atol=1e-6 after 5 SGD steps at lr=0.01. A draw where one branch has near-vanishing probability, or where the retained key is insensitive to that branch's phase, fails non-reproducibly. Add torch.manual_seed(...) and a fixed pcvl.BS()-based prefix.

Compounding it, lines 270-274:

first_probability_by_branch = {key[0]: probability for key, probability in probs.items() if key[0] in (0, 1)}

keeps only the last key per key[0], so the loss depends on dict iteration order rather than on a stated quantity. Summing over all keys of each branch is both more robust and more obviously intentional.

---

Minor / nits

• Dead code, in both copies (rst + test:129-134): branch_prob_weighted = branch.probability.unsqueeze(-1) immediately followed by branch_prob_weighted.squeeze(-1) * branch_probs_for_key is a no-op round trip — just branch.probability * branch_probs_for_key. The doc bullet "use unsqueeze(-1) or indexing to ensure ... multiplied element-wise" is therefore documenting a line that does nothing, which is more confusing than omitting it.
• x=None default is a trap (test:74, doc forward(self, x=None)): with x=None and a branch that has input_params, branch_layer(None) hits TypeError: Unsupported input types: ['NoneType'] (layer.py:1058-1068). No test covers this path. Either drop the default or raise a clear error when a branch needs inputs and x is None.
• Formatting: 16 lines carry trailing whitespace, and imports aren't isort-ordered (torch before perceval) → W293 / I001. CI only lints merlin/ (.github/workflows/ci.yml:56-57) so this won't go red, but "Code formatted (ruff format)" is unchecked in the PR description.
• Source typo still quoted verbatim: feed_forward.py:699 reads "Create you own stages" (missing "r", no final period) and the doc reproduces it. Fixing it at the source plus the two consumers is a 3-line change. Fix this →
• test:120 — circuit, _trainable_params, input_params = self.branch_configs_dict[...]: circuit is unused; the doc version correctly uses _circuit. Sync them.
• n_photons (quantum dimension) — it's a photon count, not a dimension (doc comment and test:56).
• Heading underline is 62 chars for a 58-char title; every other heading in the file matches exactly. (Retracting my round-2 sentence-case nit — the file is Title Case throughout, so the new heading is consistent.)
• merlin.algorithms.layer / merlin.core.computation_space submodule imports rather than the public from merlin import QuantumLayer used by partial_measurement.rst. Consistent with the existing example in this file, so take it or leave it.
• "expected to be solved in MerLin 0.5" will rot silently. A Jira reminder to delete the section when the new backend lands would help.
• Two different systems in one section: the error repro is 4 modes / Circuit(3) branches, the workaround is 3 modes / Circuit(2) branches. One sentence ("the workaround below uses a smaller 3-mode system for readability") saves readers from hunting for the mapping.

---

Summary. The prose is now accurate, the training-safe shape is right, and the key reconstruction is correct and general — I verified the mode-ordering invariant it relies on. Two things I'd want before merge: the doc's ValueError repro has a provider width mismatch that no test covers (item 1), and the one test whose name promises generalization beyond mode 0 still passes for the bug it was written to catch (item 3). Item 4 is the coverage gap I'd add if you want the "verified" claim to carry real weight — nothing currently checks that gradients survive set_input_state.

I did not change any files (review only), and I could not run pytest, ruff, or Sphinx in this sandbox — please rely on the CI run for green/red.
• branch PML-376-doc-ff-pm