fix(pt): omit absent charge-spin inputs for legacy scripted models

deepmd/pt/infer/deep_eval.py

@@ -830,6 +830,9 @@ def _eval_lower_strategy(
         output so the caller's extraction is unchanged.
         """
         inner = self.dp.model["Default"]
+        # Older scripted models predate this optional keyword. Omitting an
+        # absent value preserves their schema without discarding supplied data.
+        charge_spin_kwargs = {} if charge_spin is None else {"charge_spin": charge_spin}
         if self._uses_edge_schema:
             edge_schema = self._nlist_builder.build(
                 coord,
@@ -848,7 +851,7 @@ def _eval_lower_strategy(
                 edge_schema.edge_mask,
                 fparam=fparam,
                 aparam=aparam,
-                charge_spin=charge_spin,
+                **charge_spin_kwargs,
                 input_prec=coord.dtype,
             )
         else:
@@ -867,7 +870,7 @@ def _eval_lower_strategy(
                 fparam=fparam,
                 aparam=aparam,
                 do_atomic_virial=do_atomic_virial,
-                charge_spin=charge_spin,
+                **charge_spin_kwargs,
             )
             predict = communicate_extended_output(
                 model_lower,

deepmd/pt/train/wrapper.py

@@ -187,8 +187,11 @@ def forward(
             "do_atomic_virial": do_atomic_virial,
             "fparam": fparam,
             "aparam": aparam,
-            "charge_spin": charge_spin,
         }
+        # Keep older scripted model signatures usable when this optional
+        # input is absent; explicitly supplied values must still be forwarded.
+        if charge_spin is not None:
+            input_dict["charge_spin"] = charge_spin
         has_spin = getattr(self.model[task_key], "has_spin", False)
         if callable(has_spin):
             has_spin = has_spin()

source/tests/pt/model/test_nlist_backend.py

@@ -296,3 +296,22 @@ def test_nv_matches_native_nonperiodic(pt_files, name: str) -> None:
     dp_native = DeepPot(pt_files[name], nlist_backend="native")
     dp_nv = DeepPot(pt_files[name], nlist_backend="nv")
     _assert_eval_close(dp_native, dp_nv, coords, None, atype, f"{name} nv nopbc")
+
+
+@_BACKEND_MARKS["vesin"]
+def test_legacy_scripted_model_without_charge_spin():
+    """External neighbor lists must accept older scripted method schemas."""
+    from pathlib import (
+        Path,
+    )
+
+    model = Path(__file__).resolve().parents[2] / "infer/deeppot_sea.pth"
+    native = DeepPot(str(model), nlist_backend="native")
+    external = DeepPot(str(model), nlist_backend="vesin")
+    coordinates = np.array([[0.0, 0.0, 0.0, 0.9, 0.1, 0.0, -0.2, 0.85, 0.1]])
+    box = (np.eye(3) * 10.0).reshape(1, 9)
+    types = np.array([0, 1, 1])
+    expected = native.eval(coordinates, box, types)
+    actual = external.eval(coordinates, box, types)
+    for reference, result in zip(expected, actual, strict=True):
+        np.testing.assert_allclose(result, reference, rtol=1e-10, atol=1e-10)

source/tests/pt/test_wrapper.py

@@ -21,6 +21,7 @@ def __init__(self, *, fail_forward: bool = False) -> None:
         self.scale = torch.nn.Parameter(torch.ones((), device="cpu"))
         self.fail_forward = fail_forward
         self.last_requires_grad: tuple[bool, ...] | None = None
+        self.last_charge_spin: torch.Tensor | None = None
         self.last_box: torch.Tensor | None = None
 
     def forward(
@@ -33,7 +34,8 @@ def forward(
         aparam: torch.Tensor | None = None,
         charge_spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
-        del atype, do_atomic_virial, fparam, aparam, charge_spin
+        del atype, do_atomic_virial, fparam, aparam
+        self.last_charge_spin = charge_spin
         self.last_box = box
         self.last_requires_grad = tuple(
             param.requires_grad for param in self.parameters()
@@ -51,6 +53,22 @@ def forward(
         }
 
 
+class _LegacyToyModel(torch.nn.Module):
+    """Model with the scripted signature used before charge_spin was added."""
+
+    def forward(
+        self,
+        coord: torch.Tensor,
+        atype: torch.Tensor,
+        box: torch.Tensor | None = None,
+        do_atomic_virial: bool = False,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Return a simple energy without accepting a charge_spin keyword."""
+        return {"energy": coord.sum(dim=-1).sum(dim=-1, keepdim=True)}
+
+
 class _EnergyLoss(torch.nn.Module):
     def forward(
         self,
@@ -72,6 +90,35 @@ def setUp(self) -> None:
         self.coord = torch.randn(2, 5, 3, device="cpu")
         self.atype = torch.zeros(2, 5, dtype=torch.long, device="cpu")
 
+    def test_legacy_scripted_signature_without_charge_spin(self) -> None:
+        """Omitting charge_spin must work without a neighbor-list dependency."""
+        for mode in ("inference", "skip_loss", "training"):
+            with self.subTest(mode=mode):
+                model = torch.jit.script(_LegacyToyModel())
+                loss = None if mode == "inference" else _EnergyLoss()
+                wrapper = ModelWrapper(model, loss)
+                expected = model(self.coord, self.atype)
+                actual, _, _ = wrapper(
+                    self.coord, self.atype, skip_loss=mode == "skip_loss"
+                )
+                torch.testing.assert_close(actual["energy"], expected["energy"])
+
+    def test_supplied_charge_spin_is_forwarded_unchanged(self) -> None:
+        """Forward the supplied tensor through inference and loss evaluation."""
+        charge_spin = torch.tensor([[1.0, 0.0], [-1.0, 2.0]], device="cpu")
+        for mode in ("inference", "skip_loss", "training"):
+            with self.subTest(mode=mode):
+                model = _LinearToyModel()
+                loss = None if mode == "inference" else _EnergyLoss()
+                wrapper = ModelWrapper(model, loss)
+                wrapper(
+                    self.coord,
+                    self.atype,
+                    skip_loss=mode == "skip_loss",
+                    charge_spin=charge_spin,
+                )
+                self.assertIs(model.last_charge_spin, charge_spin)
+
     def test_inference_wrapper_freezes_parameters_without_changing_predictions(
         self,
     ) -> None: