fix UT

Author: caic99

deepmd/dpmodel/descriptor/dpa3.py

@@ -150,6 +150,7 @@ def __init__(
         skip_stat: bool = False,
         optim_update: bool = True,
         smooth_edge_update: bool = False,
+        use_ext_ebd: bool = False,
     ) -> None:
         self.n_dim = n_dim
         self.e_dim = e_dim
@@ -176,6 +177,7 @@ def __init__(
         self.a_compress_use_split = a_compress_use_split
         self.optim_update = optim_update
         self.smooth_edge_update = smooth_edge_update
+        self.use_ext_ebd = use_ext_ebd
 
     def __getitem__(self, key):
         if hasattr(self, key):
@@ -207,6 +209,7 @@ def serialize(self) -> dict:
             "fix_stat_std": self.fix_stat_std,
             "optim_update": self.optim_update,
             "smooth_edge_update": self.smooth_edge_update,
+            "use_ext_ebd": self.use_ext_ebd,
         }
 
     @classmethod

deepmd/pt/model/descriptor/repflows.py

@@ -462,7 +462,7 @@ def forward(
             assert mapping is not None
             node_ebd_ext = None
             nlist = torch.gather(
-                mapping,
+                mapping.reshape(nframes, -1),
                 1,
                 nlist.reshape(nframes, -1),
             ).reshape(nlist.shape)

source/tests/pt/model/test_dpa3.py

@@ -141,6 +141,99 @@ def test_consistency(
                 rtol=rtol,
                 atol=atol,
             )
+    def test_use_ext_ebd(
+        self,
+    ) -> None:
+        rtol, atol = get_tols("float32")
+
+        nf, nloc, nnei = self.nlist.shape
+        repflow0 = RepFlowArgs(
+            n_dim=20,
+            e_dim=10,
+            a_dim=8,
+            nlayers=3,
+            e_rcut=self.rcut,
+            e_rcut_smth=self.rcut_smth,
+            e_sel=nnei,
+            a_rcut=self.rcut - 0.1,
+            a_rcut_smth=self.rcut_smth,
+            a_sel=nnei - 1,
+            a_compress_rate=0,
+            a_compress_e_rate=2,
+            a_compress_use_split=False,
+            n_multi_edge_message=1,
+            axis_neuron=4,
+            update_angle=True,
+            update_style="res_residual",
+            update_residual_init="const",
+            smooth_edge_update=True,
+            use_ext_ebd=True,
+        )
+        # dpa3 with use_ext_ebd=True
+        dd0 = DescrptDPA3(
+            self.nt,
+            repflow=repflow0,
+            # kwargs for descriptor
+            exclude_types=[],
+            precision="float32",
+            use_econf_tebd=True,
+            type_map=["O", "H"],
+            seed=GLOBAL_SEED,
+            use_ext_ebd=True,
+        ).to(env.DEVICE)
+        rd0, _, _, _, _ = dd0(
+                torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
+                torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE),
+                torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
+                torch.tensor(self.mapping, dtype=int, device=env.DEVICE),
+            )
+
+        # dpa3 with use_ext_ebd=False
+        repflow1 = RepFlowArgs(
+            n_dim=20,
+            e_dim=10,
+            a_dim=8,
+            nlayers=3,
+            e_rcut=self.rcut,
+            e_rcut_smth=self.rcut_smth,
+            e_sel=nnei,
+            a_rcut=self.rcut - 0.1,
+            a_rcut_smth=self.rcut_smth,
+            a_sel=nnei - 1,
+            a_compress_rate=0,
+            a_compress_e_rate=2,
+            a_compress_use_split=False,
+            n_multi_edge_message=1,
+            axis_neuron=4,
+            update_angle=True,
+            update_style="res_residual",
+            update_residual_init="const",
+            smooth_edge_update=True,
+            use_ext_ebd=False,
+        )
+        dd1 = DescrptDPA3(
+            self.nt,
+            repflow=repflow1,
+            # kwargs for descriptor
+            exclude_types=[],
+            precision="float32",
+            use_econf_tebd=True,
+            type_map=["O", "H"],
+            seed=GLOBAL_SEED,
+            use_ext_ebd=False,
+        ).to(env.DEVICE)
+        rd1, _, _, _, _ = dd1(
+                torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE),
+                torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE),
+                torch.tensor(self.nlist, dtype=int, device=env.DEVICE),
+                torch.tensor(self.mapping, dtype=int, device=env.DEVICE),
+            )
+        np.testing.assert_allclose(
+            rd0.detach().cpu().numpy(),
+            rd1.detach().cpu().numpy(),
+            rtol=rtol,
+            atol=atol,
+        )
 
     def test_jit(
         self,