Update repflow_layer.py

Author: iProzd

deepmd/pt/model/descriptor/repflow_layer.py

@@ -1150,6 +1150,24 @@ def forward(
             n_edge = h2.shape[0]
         del a_nlist  # may be used in the future
 
+        if self.sequential_update and self.update_angle:
+            return self._forward_sequential(
+                node_ebd_ext,
+                edge_ebd,
+                h2,
+                angle_ebd,
+                nlist,
+                nlist_mask,
+                sw,
+                a_nlist_mask,
+                a_sw,
+                edge_index,
+                angle_index,
+                nb,
+                nloc,
+                n_edge,
+            )
+
         n2e_index, n_ext2e_index = edge_index[0], edge_index[1]
         n2a_index, eij2a_index, eik2a_index = (
             angle_index[0],
@@ -1166,95 +1184,6 @@ def forward(
             )
         )
 
-        # Edge self update (always from original embeddings)
-        edge_self_update = self._compute_edge_self_update(
-            node_ebd,
-            node_ebd_ext,
-            edge_ebd,
-            nei_node_ebd,
-            nlist,
-            n2e_index,
-            n_ext2e_index,
-        )
-
-        if self.sequential_update and self.update_angle:
-            # === Sequential update path ===
-            # Phase 1: Apply edge self residual
-            edge_ebd_s1 = edge_ebd + self.e_residual[0] * edge_self_update
-
-            # Phase 2: Angle self (uses updated edge_ebd_s1)
-            node_for_a, edge_for_a = self._prepare_angle_embeddings(
-                node_ebd, edge_ebd_s1, a_nlist_mask
-            )
-            angle_self_update = self._compute_angle_update(
-                angle_ebd,
-                node_for_a,
-                edge_for_a,
-                "angle",
-                n2a_index,
-                eij2a_index,
-                eik2a_index,
-            )
-            a_updated = angle_ebd + self.a_residual[0] * angle_self_update
-
-            # Phase 3: Edge angle (uses updated angle a_updated + edge_ebd_s1)
-            edge_angle_update = self._compute_angle_update(
-                a_updated,
-                node_for_a,
-                edge_for_a,
-                "edge",
-                n2a_index,
-                eij2a_index,
-                eik2a_index,
-            )
-            edge_angle_processed = self._compute_edge_angle_reduction(
-                edge_angle_update,
-                edge_ebd_s1,
-                a_sw,
-                a_nlist_mask,
-                nb,
-                nloc,
-                n_edge,
-                eij2a_index,
-            )
-            e_updated = edge_ebd_s1 + self.e_residual[1] * edge_angle_processed
-
-            # Phase 4+5: Node updates (uses e_updated)
-            node_self_mlp = self.act(self.node_self_mlp(node_ebd))
-            node_sym = self._compute_node_sym(
-                e_updated, nei_node_ebd, h2, nlist_mask, sw, n2e_index, nb, nloc
-            )
-            node_edge_update = self._compute_node_edge_message(
-                node_ebd,
-                node_ebd_ext,
-                e_updated,
-                nei_node_ebd,
-                sw,
-                nlist,
-                n2e_index,
-                n_ext2e_index,
-                nb,
-                nloc,
-            )
-
-            n_update_list: list[torch.Tensor] = [
-                node_ebd,
-                node_self_mlp,
-                node_sym,
-            ]
-            if self.n_multi_edge_message > 1:
-                node_edge_update_mul_head = node_edge_update.view(
-                    nb, nloc, self.n_multi_edge_message, self.n_dim
-                )
-                for head_index in range(self.n_multi_edge_message):
-                    n_update_list.append(node_edge_update_mul_head[..., head_index, :])
-            else:
-                n_update_list.append(node_edge_update)
-            n_updated = self.list_update(n_update_list, "node")
-
-            return n_updated, e_updated, a_updated
-
-        # === Parallel update path ===
         n_update_list: list[torch.Tensor] = [node_ebd]
         e_update_list: list[torch.Tensor] = [edge_ebd]
         a_update_list: list[torch.Tensor] = [angle_ebd]
@@ -1295,46 +1224,56 @@ def forward(
         n_updated = self.list_update(n_update_list, "node")
 
         # edge self message
+        edge_self_update = self._compute_edge_self_update(
+            node_ebd,
+            node_ebd_ext,
+            edge_ebd,
+            nei_node_ebd,
+            nlist,
+            n2e_index,
+            n_ext2e_index,
+        )
         e_update_list.append(edge_self_update)
 
         if self.update_angle:
             assert self.angle_self_linear is not None
             assert self.edge_angle_linear1 is not None
             assert self.edge_angle_linear2 is not None
 
-            node_for_a, edge_for_a = self._prepare_angle_embeddings(
+            node_ebd_for_angle, edge_ebd_for_angle = self._prepare_angle_embeddings(
                 node_ebd, edge_ebd, a_nlist_mask
             )
 
             # edge angle message
             edge_angle_update = self._compute_angle_update(
                 angle_ebd,
-                node_for_a,
-                edge_for_a,
+                node_ebd_for_angle,
+                edge_ebd_for_angle,
                 "edge",
                 n2a_index,
                 eij2a_index,
                 eik2a_index,
             )
-            edge_angle_processed = self._compute_edge_angle_reduction(
-                edge_angle_update,
-                edge_ebd,
-                a_sw,
-                a_nlist_mask,
-                nb,
-                nloc,
-                n_edge,
-                eij2a_index,
+            e_update_list.append(
+                self._compute_edge_angle_reduction(
+                    edge_angle_update,
+                    edge_ebd,
+                    a_sw,
+                    a_nlist_mask,
+                    nb,
+                    nloc,
+                    n_edge,
+                    eij2a_index,
+                )
             )
-            e_update_list.append(edge_angle_processed)
             # update edge_ebd
             e_updated = self.list_update(e_update_list, "edge")
 
             # angle self message
             angle_self_update = self._compute_angle_update(
                 angle_ebd,
-                node_for_a,
-                edge_for_a,
+                node_ebd_for_angle,
+                edge_ebd_for_angle,
                 "angle",
                 n2a_index,
                 eij2a_index,
@@ -1349,6 +1288,129 @@ def forward(
         a_updated = self.list_update(a_update_list, "angle")
         return n_updated, e_updated, a_updated
 
+    def _forward_sequential(
+        self,
+        node_ebd_ext: torch.Tensor,
+        edge_ebd: torch.Tensor,
+        h2: torch.Tensor,
+        angle_ebd: torch.Tensor,
+        nlist: torch.Tensor,
+        nlist_mask: torch.Tensor,
+        sw: torch.Tensor,
+        a_nlist_mask: torch.Tensor,
+        a_sw: torch.Tensor,
+        edge_index: torch.Tensor,
+        angle_index: torch.Tensor,
+        nb: int,
+        nloc: int,
+        n_edge: int | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Sequential update path: edge_self -> angle_self -> edge_angle -> node.
+
+        Each phase consumes the most recent update of its inputs (instead of
+        the parallel path which uses the same ``edge_ebd``/``angle_ebd`` for
+        all branches). Residuals are applied immediately after each phase.
+        """
+        assert self.update_style == "res_residual"
+        assert self.update_angle
+        assert self.angle_self_linear is not None
+        assert self.edge_angle_linear1 is not None
+        assert self.edge_angle_linear2 is not None
+        node_ebd = node_ebd_ext[:, :nloc, :]
+        n2e_index, n_ext2e_index = edge_index[0], edge_index[1]
+        n2a_index, eij2a_index, eik2a_index = (
+            angle_index[0],
+            angle_index[1],
+            angle_index[2],
+        )
+        nei_node_ebd = (
+            _make_nei_g1(node_ebd_ext, nlist)
+            if not self.use_dynamic_sel
+            else torch.index_select(
+                node_ebd_ext.reshape(-1, self.n_dim), 0, n_ext2e_index
+            )
+        )
+
+        # Phase 1: edge self update -> apply residual immediately.
+        edge_self_update = self._compute_edge_self_update(
+            node_ebd,
+            node_ebd_ext,
+            edge_ebd,
+            nei_node_ebd,
+            nlist,
+            n2e_index,
+            n_ext2e_index,
+        )
+        edge_ebd_s1 = edge_ebd + self.e_residual[0] * edge_self_update
+
+        # Phase 2: angle self message uses the updated edge embedding.
+        node_ebd_for_angle, edge_ebd_for_angle = self._prepare_angle_embeddings(
+            node_ebd, edge_ebd_s1, a_nlist_mask
+        )
+        angle_self_update = self._compute_angle_update(
+            angle_ebd,
+            node_ebd_for_angle,
+            edge_ebd_for_angle,
+            "angle",
+            n2a_index,
+            eij2a_index,
+            eik2a_index,
+        )
+        a_updated = angle_ebd + self.a_residual[0] * angle_self_update
+
+        # Phase 3: edge angle message uses the updated angle embedding.
+        edge_angle_update = self._compute_angle_update(
+            a_updated,
+            node_ebd_for_angle,
+            edge_ebd_for_angle,
+            "edge",
+            n2a_index,
+            eij2a_index,
+            eik2a_index,
+        )
+        edge_angle_reduced = self._compute_edge_angle_reduction(
+            edge_angle_update,
+            edge_ebd_s1,
+            a_sw,
+            a_nlist_mask,
+            nb,
+            nloc,
+            n_edge,
+            eij2a_index,
+        )
+        e_updated = edge_ebd_s1 + self.e_residual[1] * edge_angle_reduced
+
+        # Phase 4: node updates use the fully updated edge embedding.
+        n_update_list: list[torch.Tensor] = [node_ebd]
+        n_update_list.append(self.act(self.node_self_mlp(node_ebd)))
+        n_update_list.append(
+            self._compute_node_sym(
+                e_updated, nei_node_ebd, h2, nlist_mask, sw, n2e_index, nb, nloc
+            )
+        )
+        node_edge_update = self._compute_node_edge_message(
+            node_ebd,
+            node_ebd_ext,
+            e_updated,
+            nei_node_ebd,
+            sw,
+            nlist,
+            n2e_index,
+            n_ext2e_index,
+            nb,
+            nloc,
+        )
+        if self.n_multi_edge_message > 1:
+            node_edge_update_mul_head = node_edge_update.view(
+                nb, nloc, self.n_multi_edge_message, self.n_dim
+            )
+            for head_index in range(self.n_multi_edge_message):
+                n_update_list.append(node_edge_update_mul_head[..., head_index, :])
+        else:
+            n_update_list.append(node_edge_update)
+        n_updated = self.list_update(n_update_list, "node")
+        return n_updated, e_updated, a_updated
+
     @torch.jit.export
     def list_update_res_avg(
         self,