diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index 272e39227d..c282db4906 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/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, diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index 7d7308b0e9..faf9fe3c25 100644 --- a/deepmd/pt/train/wrapper.py +++ b/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() diff --git a/source/tests/pt/model/test_nlist_backend.py b/source/tests/pt/model/test_nlist_backend.py index 7b757efd65..b48831b2ee 100644 --- a/source/tests/pt/model/test_nlist_backend.py +++ b/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) diff --git a/source/tests/pt/test_wrapper.py b/source/tests/pt/test_wrapper.py index 5d74ff7b1b..14e8525183 100644 --- a/source/tests/pt/test_wrapper.py +++ b/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: