fix(pd): correct tensor reshaping and indexing

Author: caic99

deepmd/pd/model/descriptor/repflow_layer.py

@@ -586,7 +586,7 @@ def optim_angle_update_dynamic(
         sub_node_update = paddle.matmul(node_ebd, sub_node)
         # n_angle * angle_dim
         sub_node_update = paddle.index_select(
-            sub_node_update.reshape(nf * nloc, sub_node_update.shape[-1]), n2a_index, 0
+            sub_node_update.reshape((nf * nloc, sub_node_update.shape[-1])), n2a_index, 0
         )
 
         # n_edge * angle_dim
@@ -666,7 +666,7 @@ def optim_edge_update_dynamic(
         sub_node_update = paddle.matmul(node_ebd, node)
         # n_edge * node/edge_dim
         sub_node_update = paddle.index_select(
-            sub_node_update.reshape(nf * nloc, sub_node_update.shape[-1]),
+            sub_node_update.reshape((nf * nloc, sub_node_update.shape[-1])),
             n2e_index,
             0,
         )
@@ -675,7 +675,7 @@ def optim_edge_update_dynamic(
         sub_node_ext_update = paddle.matmul(node_ebd_ext, node_ext)
         # n_edge * node/edge_dim
         sub_node_ext_update = paddle.index_select(
-            sub_node_ext_update.reshape(nf * nall, sub_node_update.shape[-1]),
+            sub_node_ext_update.reshape((nf * nall, sub_node_update.shape[-1])),
             n_ext2e_index,
             0,
         )
@@ -746,7 +746,7 @@ def forward(
         a_updated : nf x nloc x a_nnei x a_nnei x a_dim
             Updated angle embedding.
         """
-        nb, nloc, nnei, _ = edge_ebd.shape
+        nb, nloc, nnei = nlist.shape
         nall = node_ebd_ext.shape[1]
         node_ebd = node_ebd_ext[:, :nloc, :]
         n_edge = int(nlist_mask.sum().item())
@@ -896,7 +896,7 @@ def forward(
                     n2e_index,
                     average=False,
                     num_owner=nb * nloc,
-                ).reshape(nb, nloc, node_edge_update.shape[-1])
+                ).reshape((nb, nloc, node_edge_update.shape[-1]))
                 / self.dynamic_e_sel
             )
         )

deepmd/pd/model/network/utils.py

@@ -30,7 +30,7 @@ def aggregate(
     output: [num_owner, feature_dim]
     """
     bin_count = paddle.bincount(owners)
-    bin_count = bin_count.where(bin_count != 0, paddle.ones_like(bin_count))
+    bin_count = paddle.where(bin_count != 0, bin_count, paddle.ones_like(bin_count))
 
     if (num_owner is not None) and (bin_count.shape[0] != num_owner):
         difference = num_owner - bin_count.shape[0]
@@ -51,6 +51,7 @@ def get_graph_index(
     nlist_mask: paddle.Tensor,
     a_nlist_mask: paddle.Tensor,
     nall: int,
+    use_loc_mapping: bool = True,
 ):
     """
     Get the index mapping for edge graph and angle graph, ready in `aggregate` or `index_select`.
@@ -100,7 +101,9 @@ def get_graph_index(
     n2e_index = n2e_index[nlist_mask]  # graph node index, atom_graph[:, 0]
 
     # node_ext(j) to edge(ij) index_select
-    frame_shift = paddle.arange(0, nf, dtype=nlist.dtype) * nall
+    frame_shift = paddle.arange(0, nf, dtype=nlist.dtype) * (
+        nall if not use_loc_mapping else nloc
+    )
     shifted_nlist = nlist + frame_shift[:, None, None]
     # n_edge
     n_ext2e_index = shifted_nlist[nlist_mask]  # graph neighbor index, atom_graph[:, 1]