PML-384 amps bug

[ben9871]

Summary

QCNNClassifier.amplitude_encode hardcoded torch.complex64. After model.to(torch.float64) the quantum layers correctly switch to complex128, but the encoder kept truncating every input to single precision at the pipeline entry. The layers then silently upcast the state, so the model returned float64 logits computed from float32-quality amplitudes: no error, no warning, and a double-precision QCNN produced the same numbers as a float32 one (measured ~1e-8 truncation on encoded amplitudes, ~7e-8 agreement between a float64 and a float32 model on identical data).

The fix derives the encoding dtype from the model, so model.to(torch.float64) now gives double precision end to end. The contract, pinned by regression tests: a module computes in its parameter dtype; the encoding follows the model dtype regardless of the input tensor's dtype; unsupported dtypes such as float16 raise ValueError.

This mainly affects anyone using .double() for numerical validation — torch.autograd.gradcheck requires float64 and perturbs at ~1e-6, so the ~1e-8 encoding noise could produce marginal or flaky Jacobian mismatches with no visible cause.

Related Issue

Related Jira: PML-384

Type of change

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

Proposed changes

Fix (9483f9a8)

• New QCNNClassifier._model_complex_dtype() helper: returns the first quantum layer's complex_dtype (kept in sync with .to()/.double()/.float() conversions by QuantumLayer._apply), falling back to the parameter dtype through the existing merlin.utils.dtypes.complex_dtype_for utility.
• amplitude_encode uses it for both the state tensor and the pixel amplitudes instead of hardcoded complex64; docstring states the dtype contract.
• No behavior change for default float32 models: they keep producing complex64 encodings, byte for byte.

Regression tests (ad077e94), five new tests in tests/models/test_qcnn.py:

• .to(float64) produces a complex128 encoded state and float64 logits end to end.
• A default float32 model keeps the historical complex64 encoding.
• The encoding follows the model dtype, not the input dtype: a float64 input on a float32 model still encodes as complex64.
• Float64 encoding matches an exact complex128 reference to 1e-14 and measurably differs from the old complex64 truncation — this test fails on the pre-fix code.
• .to(torch.float16) raises the existing ValueError instead of degrading silently.

API changes: none. _model_complex_dtype is private; amplitude_encode's signature is unchanged.

Out of scope — related gaps found during the investigation

Both are candidates for follow-up tickets; neither causes silent wrong numbers once this fix lands.

1. QuantumLayer ignores torch.set_default_dtype() at construction. MerlinModule.setup_device_and_dtype falls back to a hardcoded torch.float32 when no dtype= is passed (merlin/algorithms/module.py:102), so set_default_dtype(float64) before construction still yields float32 quantum layers. Core torch modules honor the global default. Fixing it is one line but changes the construction default of every layer in the library — anyone setting a global double default while relying on quantum layers staying float32 would see simulation cost roughly double, so it needs its own ticket and a release-notes entry. After this PR the behavior is at least coherent: the whole pipeline stays float32 and layer.dtype reports it honestly.

2. QCNNClassifier has no dtype/device constructor arguments. Raw QuantumLayer accepts dtype= at construction; the classifier does not, so the only precision route is construct-then-.to(), which builds at float32 and converts afterwards.

How to test / How to run

1. Regression suite (the precision test fails on the pre-fix code):

pytest tests/models/test_qcnn.py -q

Expected: 17 passed, 1 skipped. Full tests/models: 101 passed, 8 skipped.

2. Manual check of the fixed behavior:

import torch
from merlin.models import QCNNClassifier

model = QCNNClassifier((4, 4), 10).to(torch.float64)
x = torch.rand(2, 1, 4, 4, dtype=torch.float64)
print(model.amplitude_encode(x).tensor.dtype)   # complex128 (was complex64)
print(model(x).dtype)                           # float64

Screenshots / Logs (optional)

Pre-fix measurement on identical data and seeds: encoded amplitudes truncated by ~2.4e-8 (complex64 round-trip), float64-model logits within ~6.8e-8 of a float32 model. Post-fix: encoding matches an exact complex128 reference to 1e-14.

Performance considerations (optional)

None for existing float32 usage (unchanged path). Float64 models now genuinely compute in double precision downstream of the encoder, with the usual ~2x memory/compute cost of complex128 simulation — which is what .to(torch.float64) was asking for.

Documentation

• User docs updated (Sphinx)
• Examples / notebooks updated
• Docstrings updated — amplitude_encode now states the dtype contract; new helper fully documented
• Updated the API — n/a: no public API change

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 — not configured for this repo; not run

• Unit tests added/updated (pytest) — 5 regression tests

• Tests pass locally (pytest) — tests/models: 101 passed, 8 skipped

• Tests pass on GPU (pytest) — not run; the change is dtype selection only, no device-specific paths

• Test coverage not decreased significantly — increased: the encoding dtype contract had no coverage

• 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 — n/a: no new public API

• Dependencies updated (if needed) and pinned appropriately — no dependency changes

• PR description explains what changed and how to validate it

[LF-Vigneux]

@claude review

[claude]

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

---

Claude Code Review

I reviewed the diff against origin/release/0.4.1. This is a clean, well-scoped bug fix with strong regression coverage. Verdict: approve-quality; no blocking issues.

• Read PR diff and changed files
• Review qcnn.py fix
• Review regression tests
• Check repo guidelines / surrounding code
• Post review

---

Summary

The fix correctly addresses the reported bug. amplitude_encode previously hardcoded torch.complex64, truncating inputs to single precision even on a .double() model, then let the layers silently upcast — the exact "float64 logits from float32 amplitudes" failure described. Deriving the encoding dtype from the model via _model_complex_dtype() resolves this end to end, and the default float32 path is byte-for-byte unchanged.

What's good

• Correct source of truth. _model_complex_dtype() reads the first QuantumLayer.complex_dtype, which QuantumLayer._apply keeps in sync on every .to()/.double()/.float() (layer.py:1541-1544). This is the authoritative value, better than re-deriving from an input tensor.
• Robust fallback. self.layers is the torch.nn.Sequential built in build_qcnn_model with QuantumLayers as direct children (qcnn.py:540/557/574), so the isinstance loop always finds one. The complex_dtype_for(next(self.parameters()).dtype) fallback also yields the correct answer if that ever changes — belt and suspenders.
• float16 is genuinely rejected, not degraded. The ValueError in the test originates from MerlinModule.setup_device_and_dtype (module.py:103-106) during .to(torch.float16), so half precision fails at conversion time and never reaches the encoder. Worth noting: complex_dtype_for itself would happily map float16 → complex32 (dtypes.py:73-74), so the guard rail is the .to() path, not the helper — the test correctly targets .to().
• Tests are the strong part of this PR. Five focused tests pinning the contract (model-dtype not input-dtype, default stays complex64, exact complex128 reference to 1e-14, and a case that fails pre-fix). The test_amplitude_encode_float64_preserves_double_precision reference reconstruction mirrors the implementation and asserts the old truncation would deviate — a real regression guard.

Minor observations (non-blocking)

1. next(self.parameters()) in the fallback raises StopIteration if the model has no parameters. Unreachable for QCNNClassifier in practice, but next(self.parameters(), None) with a guard (or iterating self.modules()) would be marginally safer. Not worth changing on its own.
2. Pre-existing to() override (qcnn.py:497-503) — outside this PR's diff — calls super().to(...) (which already recurses into children via _apply) and then re-calls layer.to(...) on each QuantumLayer. That's redundant, though harmless. Flagging only since it's adjacent to the change; not something to fix here.
3. The out-of-scope gaps you documented (global set_default_dtype ignored at construction; no dtype=/device= ctor args on QCNNClassifier) are the right call to defer — as you note, neither produces silent wrong numbers after this lands. Good follow-up ticket candidates.

Testing note

I was unable to execute pytest tests/models/test_qcnn.py in this environment (command approval not granted), so I'm relying on code inspection plus your reported run (17 passed, 1 skipped). The tests read as correct and the pre-fix-failing assertion is sound.

Overall: correct fix, minimal blast radius, good tests. LGTM.
• PML-384_amps_bug

[LF-Vigneux]

Verything sems good! Claudes comments are out of scope