refactor: Refactor MACE wrapper helper code

deepmd_gnn/deepmd_ops.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Small compatibility hooks for DeePMD custom PyTorch ops."""
+
+import torch
+
+
+def ensure_border_op_placeholder() -> None:
+    """Install a TorchScript-visible placeholder for ``deepmd::border_op``.
+
+    DeePMD-kit provides the real op when its PyTorch extension is loaded. During
+    unit tests, documentation builds, or partial imports that extension may not
+    be present yet, but TorchScript still needs an attribute with the expected
+    name so model classes can be scripted.
+    """
+    if hasattr(torch.ops.deepmd, "border_op"):
+        return
+
+    def border_op(
+        _argument0: torch.Tensor,
+        _argument1: torch.Tensor,
+        _argument2: torch.Tensor,
+        _argument3: torch.Tensor,
+        _argument4: torch.Tensor,
+        _argument5: torch.Tensor,
+        _argument6: torch.Tensor,
+        _argument7: torch.Tensor,
+        _argument8: torch.Tensor,
+    ) -> list[torch.Tensor]:
+        """Runtime placeholder used when DeePMD's PT op is not loaded."""
+        msg = (
+            "border_op is unavailable. Build/load DeePMD-kit's PyTorch OP "
+            "library before running MPI message-passing inference."
+        )
+        raise NotImplementedError(msg)
+
+    torch.ops.deepmd.border_op = border_op

deepmd_gnn/export.py

@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Helpers used by PyTorch-exportable model paths."""
+
+import torch
+
+
+def pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor:
+    """Append a sentinel neighbor so symbolic export keeps neighbor axes dynamic."""
+    pad = -torch.ones(
+        (*nlist.shape[:2], 1),
+        dtype=nlist.dtype,
+        device=nlist.device,
+    )
+    return torch.cat([nlist, pad], dim=-1)
+
+
+def clear_export_guards_once(traced: torch.nn.Module) -> None:
+    """Clear over-specialized guards from the next export of ``traced``.
+
+    ``make_fx`` traces may specialize symbolic atom counts too aggressively.
+    DeePMD's export path calls ``torch.export.export`` immediately after tracing,
+    so this one-shot wrapper relaxes those constraints only for that export.
+    """
+    original_export = torch.export.export
+
+    def strip_deferred_assertions(exported: torch.export.ExportedProgram) -> None:
+        graph = exported.graph_module.graph
+        for node in list(graph.nodes):
+            if (
+                node.op == "call_function"
+                and node.target is torch.ops.aten._assert_scalar.default  # noqa: SLF001
+            ):
+                node.args = (True, node.args[1])
+        exported.graph_module.recompile()
+
+    def relax_range_constraints(exported: torch.export.ExportedProgram) -> None:
+        relaxed = exported.range_constraints.copy()
+        for symbol, value_range in exported.range_constraints.items():
+            try:
+                should_relax = bool(value_range.lower > 1)
+            except TypeError:
+                should_relax = False
+            if should_relax:
+                relaxed[symbol] = type(value_range)(1, value_range.upper)
+        exported._range_constraints = relaxed  # noqa: SLF001
+
+    def export_with_guard_cleanup(
+        *export_args: object,
+        **export_kwargs: object,
+    ) -> torch.export.ExportedProgram:
+        try:
+            exported = original_export(*export_args, **export_kwargs)
+            if export_args and export_args[0] is traced:
+                exported._guards_code = []  # noqa: SLF001
+                strip_deferred_assertions(exported)
+                relax_range_constraints(exported)
+            return exported
+        finally:
+            if torch.export.export is export_with_guard_cleanup:
+                torch.export.export = original_export
+
+    torch.export.export = export_with_guard_cleanup