PML-384 amps bug

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)

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)