remove node_ebd_ext

Author: caic99

deepmd/pt/model/descriptor/repflow_layer.py

@@ -466,7 +466,8 @@ def optim_edge_update(
 
     def forward(
         self,
-        node_ebd_ext: torch.Tensor,  # nf x nall x n_dim
+        node_ebd: torch.Tensor,  # nf x nloc x n_dim
+        node_ebd_ext: Optional[torch.Tensor],  # nf x nall x n_dim
         edge_ebd: torch.Tensor,  # nf x nloc x nnei x e_dim
         h2: torch.Tensor,  # nf x nloc x nnei x 3
         angle_ebd: torch.Tensor,  # nf x nloc x a_nnei x a_nnei x a_dim
@@ -511,8 +512,6 @@ def forward(
             Updated angle embedding.
         """
         nb, nloc, nnei, _ = edge_ebd.shape
-        nall = node_ebd_ext.shape[1]
-        node_ebd = node_ebd_ext[:, :nloc, :]
         assert (nb, nloc) == node_ebd.shape[:2]
         assert (nb, nloc, nnei) == h2.shape[:3]
         del a_nlist  # may be used in the future
@@ -524,8 +523,10 @@ def forward(
         # node self mlp
         node_self_mlp = self.act(self.node_self_mlp(node_ebd))
         n_update_list.append(node_self_mlp)
-
-        nei_node_ebd = _make_nei_g1(node_ebd_ext, nlist)
+        if node_ebd_ext is not None:
+            nei_node_ebd = _make_nei_g1(node_ebd_ext, nlist)
+        else:
+            nei_node_ebd = _make_nei_g1(node_ebd, nlist)
 
         # node sym (grrg + drrd)
         node_sym_list: list[torch.Tensor] = []

deepmd/pt/model/descriptor/repflows.py

@@ -461,16 +461,16 @@ def forward(
         # nb x nall x n_dim
         if comm_dict is None:
             assert mapping is not None
-            mapping = (
-                mapping.view(nframes, nall).unsqueeze(-1).expand(-1, -1, self.n_dim)
-            )
+            node_ebd_ext = None
+            nlist = torch.gather(
+                mapping,
+                1,
+                nlist.reshape(nframes, -1),
+            ).reshape(nlist.shape)
         for idx, ll in enumerate(self.layers):
             # node_ebd:     nb x nloc x n_dim
             # node_ebd_ext: nb x nall x n_dim
-            if comm_dict is None:
-                assert mapping is not None
-                node_ebd_ext = torch.gather(node_ebd, 1, mapping)
-            else:
+            if comm_dict is not None:
                 has_spin = "has_spin" in comm_dict
                 if not has_spin:
                     n_padding = nall - nloc
@@ -528,6 +528,7 @@ def forward(
                         node_ebd_real_ext, node_ebd_virtual_ext, real_nloc
                     )
             node_ebd, edge_ebd, angle_ebd = ll.forward(
+                node_ebd,
                 node_ebd_ext,
                 edge_ebd,
                 h2,