fix: DPA1 should use masked gate when excluded_types (#3815)

Fix #3707.
deepmodeling · May 24, 2024 · 8cd3cba · 8cd3cba
1 parent 0e67a7e
commit 8cd3cba
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 7 deletions.
diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py
@@ -845,10 +845,8 @@ def call(
         else:
             raise NotImplementedError
 
-        input_r = dmatrix.reshape(-1, nnei, 4)[:, :, 1:4] / np.maximum(
-            np.linalg.norm(
-                dmatrix.reshape(-1, nnei, 4)[:, :, 1:4], axis=-1, keepdims=True
-            ),
+        input_r = rr.reshape(-1, nnei, 4)[:, :, 1:4] / np.maximum(
+            np.linalg.norm(rr.reshape(-1, nnei, 4)[:, :, 1:4], axis=-1, keepdims=True),
             1e-12,
         )
         gg = self.dpa1_attention(

diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py
@@ -559,7 +559,7 @@ def forward(
                 raise NotImplementedError
 
             input_r = torch.nn.functional.normalize(
-                dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1
+                rr.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1
             )
             gg = self.dpa1_attention(
                 gg, nlist_mask, input_r=input_r, sw=sw

diff --git a/source/tests/consistent/descriptor/test_dpa1.py b/source/tests/consistent/descriptor/test_dpa1.py
@@ -201,15 +201,13 @@ def skip_tf(self) -> bool:
             precision,
             use_econf_tebd,
         ) = self.param
-        # TODO (excluded_types != [] and attn_layer > 0) need fix
         return (
             CommonTest.skip_tf
             or (
                 env_protection != 0.0
                 or smooth_type_embedding
                 or not normalize
                 or temperature != 1.0
-                or (excluded_types != [] and attn_layer > 0)
                 or (type_one_side and tebd_input_mode == "strip")  # not consistent yet
             )
             or self.is_meaningless_zero_attention_layer_tests(