From 9483f9a872f6607bc5a6c5d1e4d790ceb7aa2042 Mon Sep 17 00:00:00 2001 From: Benjamin Stott <90058728+ben9871@users.noreply.github.com> Date: Wed, 22 Jul 2026 14:46:17 +0200 Subject: [PATCH 1/2] Derive QCNN amplitude-encoding dtype from the model QCNNClassifier.amplitude_encode hardcoded torch.complex64 for the encoded state vector and pixel amplitudes. 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: a double-precision QCNN produced the same numbers as a float32 one (~1e-7 agreement) while every visible signal claimed float64. Derive the encoding dtype from the model instead: use the first quantum layer's complex_dtype (kept authoritative by QuantumLayer._apply during to()/double()/float() conversions), falling back to the parameter dtype via the existing complex_dtype_for utility. The encoding follows the model dtype regardless of the input tensor's dtype, matching the standard torch contract that a module computes in its parameter dtype. --- merlin/models/qcnn.py | 31 +++++++++++++++++++++++++++++-- 1 file changed, 29 insertions(+), 2 deletions(-) diff --git a/merlin/models/qcnn.py b/merlin/models/qcnn.py index 76ea4707b..dc7beeeb2 100644 --- a/merlin/models/qcnn.py +++ b/merlin/models/qcnn.py @@ -13,6 +13,7 @@ from ..core import ComputationSpace, StateVector from ..core.partial_measurement import PartialMeasurement from ..measurement import MeasurementStrategy +from ..utils.dtypes import complex_dtype_for class QCNNClassifier(torch.nn.Module): @@ -897,6 +898,26 @@ def _build_amplitude_basis_indices(self) -> torch.Tensor: return torch.tensor(basis_indices, dtype=torch.long) + def _model_complex_dtype(self) -> torch.dtype: + """Return the complex dtype matching the model's floating-point dtype. + + The quantum layers hold the authoritative ``complex_dtype`` (kept in + sync with device/dtype conversions such as + :meth:`~torch.nn.Module.to` and :meth:`~torch.nn.Module.double`), so + the first quantum layer's value is used. If no quantum layer exists, + the dtype is derived from the model parameters instead. + + Returns + ------- + torch.dtype + ``torch.complex64`` for a float32 model, ``torch.complex128`` + for a float64 model. + """ + for layer in self.layers: + if isinstance(layer, QuantumLayer): + return layer.complex_dtype + return complex_dtype_for(next(self.parameters()).dtype) + def amplitude_encode(self, x: torch.Tensor): """Encode image pixels as amplitudes of a two-photon state vector. @@ -905,6 +926,11 @@ def amplitude_encode(self, x: torch.Tensor): ``input_shape[0] + j`` of the second register. The resulting state vector is normalized before being passed to the quantum layers. + The encoding follows the model's precision: amplitudes are created + with the complex dtype matching the model dtype (``complex64`` for a + float32 model, ``complex128`` for a float64 model), regardless of the + input tensor's dtype. + Parameters ---------- x : torch.Tensor @@ -924,10 +950,11 @@ def amplitude_encode(self, x: torch.Tensor): ) basis_size = state_vector.basis_size + complex_dtype = self._model_complex_dtype() state_tensor = torch.zeros( - (batch_size, basis_size), dtype=torch.complex64, device=x.device + (batch_size, basis_size), dtype=complex_dtype, device=x.device ) - pixel_amplitudes = x[:, 0, :, :].reshape(batch_size, -1).to(torch.complex64) + pixel_amplitudes = x[:, 0, :, :].reshape(batch_size, -1).to(complex_dtype) basis_indices = self._amplitude_basis_indices.to(x.device) expanded_indices = basis_indices.unsqueeze(0).expand(batch_size, -1) state_tensor = state_tensor.scatter(1, expanded_indices, pixel_amplitudes) From ad077e94b8b43860f8fc42c297f625979092cb27 Mon Sep 17 00:00:00 2001 From: Benjamin Stott <90058728+ben9871@users.noreply.github.com> Date: Wed, 22 Jul 2026 14:48:02 +0200 Subject: [PATCH 2/2] Add regression tests for QCNN encoding precision Pin the amplitude-encoding dtype contract fixed in the previous commit: - to(float64) produces a complex128 encoded state and float64 logits. - 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 beats the old complex64 truncation (~1e-8), proving double precision now actually flows through the entry point. - to(torch.float16) raises the existing clear ValueError rather than degrading silently. Also removes two pre-existing whitespace lint findings picked up by ruff --fix in this file. --- tests/models/test_qcnn.py | 58 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 57 insertions(+), 1 deletion(-) diff --git a/tests/models/test_qcnn.py b/tests/models/test_qcnn.py index b9e75b3d9..5f9835ca7 100644 --- a/tests/models/test_qcnn.py +++ b/tests/models/test_qcnn.py @@ -102,7 +102,7 @@ def test_qcnn_basic_api(): ) with pytest.raises(AttributeError): qcnn_classifier_from_list.input_shape = (2, 2) - + qcnn_classifier_above_former_limit = QCNNClassifier( input_shape_above_former_limit, accepted_num_classes, @@ -791,3 +791,59 @@ def test_qcnn_state_dict_round_trip_with_export_config(): assert torch.allclose(restored_logits, expected_logits, atol=1e-6, rtol=1e-6) + +def test_amplitude_encode_follows_model_dtype_after_to_float64(): + """to(float64) yields complex128 encoding and float64 logits end to end.""" + model = QCNNClassifier((4, 4), 10).to(torch.float64) + x = torch.rand(2, 1, 4, 4, dtype=torch.float64) + + encoded = model.amplitude_encode(x) + assert encoded.tensor.dtype == torch.complex128 + + logits = model(x) + assert logits.dtype == torch.float64 + + +def test_amplitude_encode_default_model_uses_complex64(): + """A default (float32) model keeps the historical complex64 encoding.""" + model = QCNNClassifier((4, 4), 10) + x = torch.rand(2, 1, 4, 4) + + assert model.amplitude_encode(x).tensor.dtype == torch.complex64 + + +def test_amplitude_encode_input_dtype_follows_model_not_input(): + """A float64 input on a float32 model encodes at the model's precision.""" + model = QCNNClassifier((4, 4), 10) + x = torch.rand(2, 1, 4, 4, dtype=torch.float64) + + assert model.amplitude_encode(x).tensor.dtype == torch.complex64 + + +def test_amplitude_encode_float64_preserves_double_precision(): + """Encoded amplitudes match an exact complex128 reference (no truncation).""" + model = QCNNClassifier((4, 4), 10).to(torch.float64) + x = torch.rand(2, 1, 4, 4, dtype=torch.float64) + batch_size = x.shape[0] + + encoded = model.amplitude_encode(x).tensor + + pixels = x[:, 0].reshape(batch_size, -1).to(torch.complex128) + indices = model._amplitude_basis_indices.unsqueeze(0).expand(batch_size, -1) + reference = torch.zeros_like(encoded).scatter(1, indices, pixels) + reference = reference / reference.abs().pow(2).sum(dim=1, keepdim=True).sqrt() + + torch.testing.assert_close(encoded, reference, rtol=0.0, atol=1e-14) + + # The truncated complex64 encoding would deviate by ~1e-8; double precision + # must beat that by orders of magnitude. + truncated = reference.to(torch.complex64).to(torch.complex128) + assert not torch.allclose(encoded, truncated, rtol=0.0, atol=1e-12) + + +def test_to_unsupported_dtype_raises_clear_error(): + """Half precision is rejected loudly instead of degrading silently.""" + model = QCNNClassifier((4, 4), 10) + + with pytest.raises(ValueError, match="torch.float32 or torch.float64"): + model.to(torch.float16)