refine more code

Author: HydrogenSulfate

deepmd/pd/model/descriptor/dpa2.py

@@ -25,6 +25,7 @@
     TypeEmbedNetConsistent,
 )
 from deepmd.pd.utils import (
+    decomp,
     env,
 )
 from deepmd.pd.utils.env import (
@@ -799,7 +800,8 @@ def forward(
             g1 = g1 + self.tebd_transform(g1_inp)
         # mapping g1
         if comm_dict is None or len(comm_dict) == 0:
-            assert mapping.numel() > 0
+            if paddle.in_dynamic_mode():
+                assert decomp.numel(mapping) > 0
             mapping_ext = (
                 mapping.reshape([nframes, nall])
                 .unsqueeze(-1)

deepmd/pd/model/descriptor/repflows.py

@@ -5,7 +5,6 @@
     Union,
 )
 
-import numpy as np
 import paddle
 
 from deepmd.dpmodel.utils.seed import (
@@ -525,7 +524,7 @@ def forward(
         # nf x nloc x a_nnei x a_nnei
         # 1 - 1e-6 for paddle.acos stability
         cosine_ij = paddle.matmul(normalized_diff_i, normalized_diff_j) * (1 - 1e-6)
-        angle_input = cosine_ij.unsqueeze(-1) / (np.pi**0.5)
+        angle_input = cosine_ij.unsqueeze(-1) / (paddle.pi**0.5)
 
         if not parallel_mode and self.use_loc_mapping:
             if paddle.in_dynamic_mode():
@@ -597,19 +596,21 @@ def forward(
                 has_spin = len(comm_dict) >= 7
                 if not has_spin:
                     n_padding = nall - nloc
-                    # node_ebd = paddle.nn.functional.pad(
-                    #     node_ebd.squeeze(0), [0, 0, 0, n_padding], value=0.0
-                    # )
+                    if paddle.in_dynamic_mode():
+                        node_ebd = paddle.nn.functional.pad(
+                            node_ebd.squeeze(0), [0, 0, 0, n_padding], value=0.0
+                        )
+                    else:
+                        _fill_shape = node_ebd.shape[1:]
+                        _fill_shape[1] = n_padding
+                        node_ebd = paddle.concat(
+                            [
+                                node_ebd.squeeze(0),
+                                paddle.zeros(_fill_shape, dtype=node_ebd.dtype),
+                            ],
+                            axis=1,
+                        )
                     # [nframes, nloc, tebd_dim]
-                    _shapes = node_ebd.shape[1:]
-                    _shapes[1] = n_padding
-                    node_ebd = paddle.concat(
-                        [
-                            node_ebd.squeeze(0),
-                            paddle.zeros(_shapes, dtype=node_ebd.dtype),
-                        ],
-                        axis=1,
-                    )
                     real_nloc = nloc
                     real_nall = nall
                 else:
@@ -630,12 +631,6 @@ def forward(
                     )
 
                 assert len(comm_dict) >= 6
-                # assert "send_list" in comm_dict
-                # assert "send_proc" in comm_dict
-                # assert "recv_proc" in comm_dict
-                # assert "send_num" in comm_dict
-                # assert "recv_num" in comm_dict
-                # assert "communicator" in comm_dict
                 ret = paddle_ops_deepmd_border_op(
                     comm_dict[0],
                     comm_dict[1],

deepmd/pd/model/descriptor/repformer_layer.py

@@ -110,7 +110,6 @@ def _make_nei_g1(
     # index: nb x (nloc x nnei) x ng1
     index = nlist.reshape([nb, nloc * nnei]).unsqueeze(-1).expand([-1, -1, ng1])
     # gg1  : nb x (nloc x nnei) x ng1
-    # print(g1_ext.shape, index.shape)
     gg1 = paddle.take_along_axis(g1_ext, indices=index, axis=1, broadcast=False)
     # gg1  : nb x nloc x nnei x ng1
     gg1 = gg1.reshape([nb, nloc, nnei, ng1])

deepmd/pd/model/descriptor/repformers.py

@@ -484,15 +484,17 @@ def forward(
                 has_spin = len(comm_dict) >= 7
                 if not has_spin:
                     n_padding = nall - nloc
-                    # g1 = paddle.nn.functional.pad(
-                    #     g1.squeeze(0), (0, 0, 0, n_padding), value=0.0
-                    # )
-                    _shapes = g1.shape[1:]
-                    _shapes[1] = n_padding
-                    g1 = paddle.concat(
-                        [g1.squeeze(0), paddle.zeros(_shapes, dtype=g1.dtype)],
-                        axis=1,
-                    )
+                    if paddle.in_dynamic_mode():
+                        g1 = paddle.nn.functional.pad(
+                            g1.squeeze(0), (0, 0, 0, n_padding), value=0.0
+                        )
+                    else:
+                        _fill_shape = g1.shape[1:]
+                        _fill_shape[1] = n_padding
+                        g1 = paddle.concat(
+                            [g1.squeeze(0), paddle.zeros(_fill_shape, dtype=g1.dtype)],
+                            axis=1,
+                        )
                     real_nloc = nloc
                     real_nall = nall
                 else:
@@ -510,13 +512,7 @@ def forward(
                         mix_g1.squeeze(0), (0, 0, 0, real_n_padding), value=0.0
                     )
 
-                # assert "send_list" in comm_dict
-                # assert "send_proc" in comm_dict
-                # assert "recv_proc" in comm_dict
-                # assert "send_num" in comm_dict
-                # assert "recv_num" in comm_dict
-                # assert "communicator" in comm_dict
-                # print(f"g1.shape = ", g1.shape)
+                assert len(comm_dict) >= 6
                 ret = paddle_ops_deepmd_border_op(
                     comm_dict[0],
                     comm_dict[1],