From cb6447ff6678be75587b917cb2b92d25c1784d2d Mon Sep 17 00:00:00 2001
From: Pascal Salzbrenner <pascal.salzbrenner@microsoft.com>
Date: Wed, 21 Oct 2026 15:03:43 +0100
Subject: [PATCH] fix: Enforce contiguous storage when compiling for AOTI

---
 src/mattersim/forcefield/aoti_compile.py | 36 +++++++++++-------
 tests/forcefield/test_aoti_compile.py    | 47 ++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 13 deletions(-)

diff --git a/src/mattersim/forcefield/aoti_compile.py b/src/mattersim/forcefield/aoti_compile.py
index a9f4a1a..e710fb4 100644
--- a/src/mattersim/forcefield/aoti_compile.py
+++ b/src/mattersim/forcefield/aoti_compile.py
@@ -474,19 +474,29 @@ def forward(
         input: dict[str, torch.Tensor],
         dataset_idx: int = -1,
     ) -> dict[str, torch.Tensor]:
-        args = (
-            input["atom_pos"],
-            input["cell"],
-            input["pbc_offsets"],
-            input["atom_attr"],
-            input["edge_index"],
-            input["three_body_indices"],
-            input["num_three_body"],
-            input["num_bonds"],
-            input["num_triple_ij"],
-            input["num_atoms"],
-            input["num_graphs"],
-            input["batch"],
+        # The AOTI kernels are specialized to the contiguous strides of the
+        # compile-time example inputs and, unless AOTI_RUNTIME_CHECK_INPUTS is
+        # set, do not check strides at runtime.
+        # Non-contiguous inputs are silently misread: e.g. TorchSim's cell
+        # filters leave ``row_vector_cell`` as a transposed view, which gave
+        # stresses off by up to thousands of GPa. ``contiguous()`` is a no-op
+        # for inputs that already have the expected layout.
+        args = tuple(
+            input[key].contiguous()
+            for key in (
+                "atom_pos",
+                "cell",
+                "pbc_offsets",
+                "atom_attr",
+                "edge_index",
+                "three_body_indices",
+                "num_three_body",
+                "num_bonds",
+                "num_triple_ij",
+                "num_atoms",
+                "num_graphs",
+                "batch",
+            )
         )
         result = self._aoti_model(*args)
         output = {
diff --git a/tests/forcefield/test_aoti_compile.py b/tests/forcefield/test_aoti_compile.py
index e33871d..6dc9b35 100644
--- a/tests/forcefield/test_aoti_compile.py
+++ b/tests/forcefield/test_aoti_compile.py
@@ -21,6 +21,7 @@
     _make_fx,
     _model_fingerprint,
 )
+from mattersim.forcefield.potential import batch_to_dict
 
 requires_gpu = pytest.mark.skipif(
     not torch.cuda.is_available(), reason="CUDA not available"
@@ -250,3 +251,49 @@ def test_energies_forces_stresses_match_eager(
         assert abs(e_ref[0] - e_aoti[0]) / len(atoms) < 1e-4
         np.testing.assert_allclose(f_aoti[0], f_ref[0], atol=1e-3)
         np.testing.assert_allclose(s_aoti[0], s_ref[0], atol=1e-2)
+
+    @pytest.mark.parametrize("cell_layout", ["contiguous", "transposed_view"])
+    def test_non_contiguous_cell_matches_eager(
+        self,
+        mattersim_potential_aoti,
+        mattersim_potential_best_device,
+        si_diamond_cubic,
+        perturb,
+        cell_layout,
+    ):
+        """Regression: AOTI misread a non-contiguous ``cell`` input.
+
+        TorchSim's cell filters store the column-vector cell, so the
+        ``row_vector_cell`` handed to the model is a transposed view with
+        strides ``(9, 1, 3)``. The artifact assumed the contiguous strides of
+        its compile-time example and returned stresses off by up to thousands
+        of GPa. A sheared (non-symmetric) cell is required: the transpose of
+        a symmetric cell has the same values in either layout.
+        """
+        atoms = perturb(si_diamond_cubic, displacement=0.05)
+        cell = atoms.cell.array.copy()
+        cell[1, 0] += 0.8
+        cell[2, 0] += 0.4
+        cell[2, 1] += 0.6
+        atoms.set_cell(cell, scale_atoms=True)
+
+        results = {}
+        for name, potential in (
+            ("eager", mattersim_potential_best_device),
+            ("aoti", mattersim_potential_aoti),
+        ):
+            graph_batch = next(iter(self._dataloader(atoms, potential)))
+            input_dict = batch_to_dict(graph_batch, device="cuda")
+            if name == "aoti" and cell_layout == "transposed_view":
+                input_dict["cell"] = input_dict["cell"].mT.contiguous().mT
+                assert not input_dict["cell"].is_contiguous()
+            output = potential(input_dict, include_forces=True, include_stresses=True)
+            results[name] = {
+                key: output[key].detach().cpu().numpy()
+                for key in ("total_energy", "forces", "stresses")
+            }
+
+        ref, aoti = results["eager"], results["aoti"]
+        assert abs(ref["total_energy"][0] - aoti["total_energy"][0]) / len(atoms) < 1e-4
+        np.testing.assert_allclose(aoti["forces"], ref["forces"], atol=1e-3)
+        np.testing.assert_allclose(aoti["stresses"], ref["stresses"], atol=1e-2)