Keep UMA atom and edge dimensions dynamic

src/fairchem/core/units/mlip_unit/predict.py

@@ -96,6 +96,26 @@ def _prepare_inference_gradients(backbone, data: AtomicData) -> None:
         data["cell"].requires_grad_(True)
 
 
+def _mark_dynamic_input_dimensions(backbone, data: AtomicData) -> None:
+    if not getattr(
+        getattr(backbone, "backend", None), "supports_fused_edgewise", False
+    ):
+        return
+    dynamic_dims = {
+        "atomic_numbers": 0,
+        "batch": 0,
+        "cell_offsets": 0,
+        "edge_index": 1,
+        "fixed": 0,
+        "pos": 0,
+        "tags": 0,
+    }
+    for key, dim in dynamic_dims.items():
+        value = data.get(key, None)
+        if torch.is_tensor(value):
+            torch._dynamo.mark_dynamic(value, dim)
+
+
 class MLIPPredictUnitProtocol(Protocol):
     def predict(self, data: AtomicData, undo_element_references: bool) -> dict: ...
 
@@ -461,6 +481,8 @@ def predict(
 
         backbone = self.model.module.backbone
         _prepare_inference_gradients(backbone, data_device)
+        if self.inference_settings.compile:
+            _mark_dynamic_input_dimensions(backbone, data_device)
 
         # Model handles any per-prediction checks (e.g., MOLE consistency)
         self.model.module.on_predict_check(data_device)
@@ -482,15 +504,14 @@ def _lazy_init(self, data: AtomicData) -> None:
             logging.warning(
                 "Model is being compiled this might take a while for the first time"
             )
-            torch._dynamo.config.recompile_limit = 32
-            # Bake float literals in as constants rather than symbolic floats.
-            # The model's scalars are fixed at inference, so this skips dynamo's
-            # TensorifyScalarRestartAnalysis retrace during compile.
-            torch._dynamo.config.specialize_float = True
             self.model = torch.compile(self.model, dynamic=True)
 
         self.lazy_model_intialized = True
 
+    @torch._dynamo.config.patch(
+        recompile_limit=32,
+        specialize_float=True,
+    )
     def _run_inference(self, data: AtomicData, undo_refs: bool) -> dict:
         """
         Execute model inference.

tests/core/units/mlip_unit/test_predict.py

@@ -44,6 +44,7 @@
 from fairchem.core.units.mlip_unit.mlip_unit import initialize_finetuning_model
 from fairchem.core.units.mlip_unit.predict import (
     ParallelMLIPPredictUnit,
+    _mark_dynamic_input_dimensions,
     _prepare_inference_gradients,
 )
 from fairchem.core.units.mlip_unit.single_atom_patch import (
@@ -95,6 +96,25 @@ def test_prepare_inference_gradients(forces, stress, pos_grad, cell_grad):
     _prepare_inference_gradients(SimpleNamespace(), data)
 
 
+def test_mark_dynamic_input_dimensions_requires_fast_backend():
+    def make_data():
+        return {
+            "pos": torch.randn(4, 3),
+            "edge_index": torch.zeros(2, 8, dtype=torch.long),
+        }
+
+    data = make_data()
+    _mark_dynamic_input_dimensions(SimpleNamespace(), data)
+    assert not hasattr(data["pos"], "_dynamo_dynamic_indices")
+    assert not hasattr(data["edge_index"], "_dynamo_dynamic_indices")
+
+    backend = SimpleNamespace(supports_fused_edgewise=True)
+    data = make_data()
+    _mark_dynamic_input_dimensions(SimpleNamespace(backend=backend), data)
+    assert data["pos"]._dynamo_dynamic_indices == {0}
+    assert data["edge_index"]._dynamo_dynamic_indices == {1}
+
+
 _REPRESENTATIVE_ELEMENTS = [
     (1, 0, 2),  # H:  charge=0, spin=2
     (6, 0, 3),  # C:  charge=0, spin=3