fix ut

Author: iProzd

deepmd/pt/modifier/base_modifier.py

@@ -83,6 +83,7 @@ def forward(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        charge_spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         """Compute energy, force, and virial corrections."""
 

deepmd/pt_expt/infer/deep_eval.py

@@ -1137,7 +1137,7 @@ def _eval_model(
                 mapping_t,
                 fparam_t,
                 aparam_t,
-                charge_spin=charge_spin_t,
+                charge_spin_t,
             )
 
         # Apply communicate_extended_output to map extended atoms → local atoms
@@ -1317,7 +1317,7 @@ def _eval_model_spin(
                 mapping_t,
                 fparam_t,
                 aparam_t,
-                charge_spin=charge_spin_t,
+                charge_spin_t,
             )
 
         # Apply communicate_extended_output to map extended atoms → local atoms

source/tests/pt/test_data_modifier.py

@@ -116,6 +116,7 @@ def forward(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        charge_spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         """Implementation of abstractmethod."""
         return {}
@@ -158,6 +159,7 @@ def forward(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        charge_spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         """Implementation of abstractmethod."""
         return {}
@@ -205,6 +207,7 @@ def forward(
         fparam: torch.Tensor | None = None,
         aparam: torch.Tensor | None = None,
         do_atomic_virial: bool = False,
+        charge_spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         """Take scaled model prediction as data modification."""
         model_pred = self.model(

source/tests/pt_expt/export_helpers.py

@@ -118,6 +118,7 @@ def model_forward_lower_export_round_trip(
     fparam,
     aparam,
     output_keys: tuple[str, ...],
+    charge_spin=None,
     rtol: float = 1e-10,
     atol: float = 1e-10,
 ):
@@ -141,6 +142,9 @@ def model_forward_lower_export_round_trip(
         Frame and atom parameters.
     output_keys : tuple of str
         Output dictionary keys to verify.
+    charge_spin : torch.Tensor or None
+        Charge/spin parameter for descriptors that consume it (e.g. DPA3
+        with ``add_chg_spin_ebd=True``).
     rtol, atol : float
         Tolerances for np.testing.assert_allclose.
     """
@@ -156,6 +160,7 @@ def model_forward_lower_export_round_trip(
         mapping_t,
         fparam=fparam,
         aparam=aparam,
+        charge_spin=charge_spin,
     )
 
     # 2. Concrete trace
@@ -166,21 +171,24 @@ def model_forward_lower_export_round_trip(
         mapping_t,
         fparam=fparam,
         aparam=aparam,
+        charge_spin=charge_spin,
     )
     assert isinstance(traced, torch.nn.Module)
 
     # 3. Basic export (no dynamic shapes)
     exported = torch.export.export(
         traced,
-        (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam),
+        (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin),
         strict=False,
     )
     assert exported is not None
 
     # 4. Compare traced and exported vs eager
-    ret_traced = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam)
+    ret_traced = traced(
+        ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin
+    )
     ret_exported = exported.module()(
-        ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam
+        ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin
     )
     for key in output_keys:
         np.testing.assert_allclose(
@@ -201,7 +209,15 @@ def model_forward_lower_export_round_trip(
     # 5. Symbolic trace + dynamic shapes + .pte round-trip
     inputs_2f = tuple(
         torch.cat([t, t], dim=0) if t is not None else None
-        for t in (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam)
+        for t in (
+            ext_coord,
+            ext_atype,
+            nlist_t,
+            mapping_t,
+            fparam,
+            aparam,
+            charge_spin,
+        )
     )
     traced_sym = md_pt.forward_lower_exportable(
         inputs_2f[0],
@@ -210,6 +226,7 @@ def model_forward_lower_export_round_trip(
         inputs_2f[3],
         fparam=inputs_2f[4],
         aparam=inputs_2f[5],
+        charge_spin=inputs_2f[6],
         tracing_mode="symbolic",
         _allow_non_fake_inputs=True,
     )
@@ -226,7 +243,9 @@ def model_forward_lower_export_round_trip(
         loaded = torch.export.load(f.name).module()
 
     # 6. Compare loaded vs eager (nf=1 — different shapes)
-    ret_loaded_1f = loaded(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam)
+    ret_loaded_1f = loaded(
+        ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin
+    )
     for key in output_keys:
         np.testing.assert_allclose(
             ret_eager[key].detach().cpu().numpy(),

source/tests/pt_expt/model/test_export_pipeline.py

@@ -121,7 +121,9 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             inputs_trace = _make_sample_inputs(model2, nframes=5, nloc=7)
         finally:
             _env.DEVICE = orig_device
-        ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam = inputs_trace
+        ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin = (
+            inputs_trace
+        )
 
         # 4. Eager reference
         eager_out = model2.forward_common_lower(
@@ -132,6 +134,7 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             fparam=fparam,
             aparam=aparam,
             do_atomic_virial=True,
+            charge_spin=charge_spin,
         )
 
         # 5. Trace with symbolic mode (same as dp freeze)
@@ -142,6 +145,7 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             mapping_t,
             fparam=fparam,
             aparam=aparam,
+            charge_spin=charge_spin,
             do_atomic_virial=True,
             tracing_mode="symbolic",
             _allow_non_fake_inputs=True,
@@ -155,11 +159,12 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             mapping_t,
             fparam,
             aparam,
+            charge_spin,
             model_nnei=sum(model2.get_sel()),
         )
         exported = torch.export.export(
             traced,
-            (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam),
+            (ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin),
             dynamic_shapes=dynamic_shapes,
             strict=False,
             prefer_deferred_runtime_asserts_over_guards=True,
@@ -171,7 +176,9 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             loaded = torch.export.load(tmp.name).module()
 
         # 8. Verify: traced output matches eager (same shapes as trace)
-        traced_out = traced(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam)
+        traced_out = traced(
+            ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin
+        )
         for key in eager_out:
             np.testing.assert_allclose(
                 eager_out[key].detach().cpu().numpy(),
@@ -182,7 +189,9 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             )
 
         # 9. Verify: loaded (.pte) output matches eager (same shapes)
-        loaded_out = loaded(ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam)
+        loaded_out = loaded(
+            ext_coord, ext_atype, nlist_t, mapping_t, fparam, aparam, charge_spin
+        )
         for key in eager_out:
             np.testing.assert_allclose(
                 eager_out[key].detach().cpu().numpy(),
@@ -206,6 +215,7 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             mapping_t2,
             fparam2,
             aparam2,
+            charge_spin2,
         ) = inputs_infer
 
         eager_out2 = model2.forward_common_lower(
@@ -216,9 +226,16 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
             fparam=fparam2,
             aparam=aparam2,
             do_atomic_virial=True,
+            charge_spin=charge_spin2,
         )
         loaded_out2 = loaded(
-            ext_coord2, ext_atype2, nlist_t2, mapping_t2, fparam2, aparam2
+            ext_coord2,
+            ext_atype2,
+            nlist_t2,
+            mapping_t2,
+            fparam2,
+            aparam2,
+            charge_spin2,
         )
         for key in eager_out2:
             np.testing.assert_allclose(
@@ -248,9 +265,16 @@ def test_export_pipeline(self, descriptor_type, with_fparam) -> None:
                 fparam=fparam_ones,
                 aparam=aparam,
                 do_atomic_virial=True,
+                charge_spin=charge_spin,
             )
             loaded_out_fp1 = loaded(
-                ext_coord, ext_atype, nlist_t, mapping_t, fparam_ones, aparam
+                ext_coord,
+                ext_atype,
+                nlist_t,
+                mapping_t,
+                fparam_ones,
+                aparam,
+                charge_spin,
             )
             # Loaded with fparam=1 should match eager with fparam=1
             for key in eager_out_fp1: