diff --git a/merlin/models/qcnn.py b/merlin/models/qcnn.py index 76ea4707..dc7beeeb 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) diff --git a/tests/models/test_qcnn.py b/tests/models/test_qcnn.py index b9e75b3d..5f9835ca 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)