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

@@ -630,12 +630,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

@@ -510,13 +510,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],