From 3c76eddcc2da7026518917ecea341cc9ffe389dc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrique=20Silv=C3=A9rio?=
 <29920212+HGSilveri@users.noreply.github.com>
Date: Thu, 31 Oct 2024 14:01:25 +0100
Subject: [PATCH] Hotfix: Preserve qubit IDs and order in
 Register.with_automatic_layout() (#759)

* Preserve qubit IDs and order in Register.with_automatic_layout()

* Bump version to v1.1.1

* Fail for registers with differentiable coordinates
---
 VERSION.txt                             |  2 +-
 pulser-core/pulser/register/register.py | 20 ++++++++++++++++++--
 tests/test_register.py                  | 15 ++++++++++++++-
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/VERSION.txt b/VERSION.txt
index 9084fa2f..524cb552 100644
--- a/VERSION.txt
+++ b/VERSION.txt
@@ -1 +1 @@
-1.1.0
+1.1.1
diff --git a/pulser-core/pulser/register/register.py b/pulser-core/pulser/register/register.py
index b110ca26..0f0be2fd 100644
--- a/pulser-core/pulser/register/register.py
+++ b/pulser-core/pulser/register/register.py
@@ -344,6 +344,8 @@ def with_automatic_layout(
         Raises:
             RuntimeError: If the automatic layout generation fails to meet
                 the device constraints.
+            NotImplementedError: When the register has differentiable
+                coordinates (ie torch Tensors with requires_grad=True).
 
         Returns:
             Register: A new register instance with identical qubit IDs and
@@ -353,6 +355,15 @@ def with_automatic_layout(
             raise TypeError(
                 f"'device' must be of type Device, not {type(device)}."
             )
+        if (
+            self._coords_arr.is_tensor
+            and self._coords_arr.as_tensor().requires_grad
+        ):
+            raise NotImplementedError(
+                "'Register.with_automatic_layout()' does not support "
+                "registers with differentiable coordinates."
+            )
+
         trap_coords = generate_trap_coordinates(
             self.sorted_coords,
             min_trap_dist=device.min_atom_distance,
@@ -363,8 +374,13 @@ def with_automatic_layout(
             max_traps=device.max_layout_traps,
         )
         layout = pulser.register.RegisterLayout(trap_coords, slug=layout_slug)
-        trap_ids = layout.get_traps_from_coordinates(*self.sorted_coords)
-        return cast(Register, layout.define_register(*trap_ids))
+        trap_ids = layout.get_traps_from_coordinates(
+            *self._coords_arr.as_array()
+        )
+        return cast(
+            Register,
+            layout.define_register(*trap_ids, qubit_ids=self.qubit_ids),
+        )
 
     def rotated(self, degrees: float) -> Register:
         """Makes a new rotated register.
diff --git a/tests/test_register.py b/tests/test_register.py
index a782bff1..5cbacf0a 100644
--- a/tests/test_register.py
+++ b/tests/test_register.py
@@ -593,7 +593,7 @@ def test_register_recipes_torch(
 
 @pytest.mark.parametrize("optimal_filling", [None, 0.4, 0.1])
 def test_automatic_layout(optimal_filling):
-    reg = Register.square(4, spacing=5)
+    reg = Register.square(4, spacing=5, prefix="test")
     max_layout_filling = 0.5
     min_traps = int(np.ceil(len(reg.qubits) / max_layout_filling))
     optimal_traps = int(
@@ -610,6 +610,8 @@ def test_automatic_layout(optimal_filling):
 
     # On its own, it works
     new_reg = reg.with_automatic_layout(device, layout_slug="foo")
+    assert new_reg.qubit_ids == reg.qubit_ids  # Same IDs in the same order
+    assert new_reg == reg  # The register itself is identical
     assert isinstance(new_reg.layout, RegisterLayout)
     assert str(new_reg.layout) == "foo"
     trap_num = new_reg.layout.number_of_traps
@@ -657,3 +659,14 @@ def test_automatic_layout(optimal_filling):
         big_reg.with_automatic_layout(device).layout.number_of_traps
         >= min_traps
     )
+
+
+def test_automatic_layout_diff():
+    torch = pytest.importorskip("torch")
+    with pytest.raises(
+        NotImplementedError,
+        match="does not support registers with differentiable coordinates",
+    ):
+        Register.square(
+            2, spacing=torch.tensor(10.0, requires_grad=True)
+        ).with_automatic_layout(AnalogDevice)