pd: Ignore if branch of 0-size (#4617)

The support for higher-order differentiation and complex control flow in
`paddle.jit` static graphs is not very strong. Therefore, branches
related to 0-size are ignored to avoid accuracy issues during training.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Bug Fixes**
- Improved the consistency and stability of key computational processes
across the application, ensuring that data processing occurs reliably
under all conditions.
- Enhanced compatibility with the underlying execution framework,
leading to more predictable outcomes even in edge-case scenarios.
- Adjusted control flow for handling empty coordinate tensors, allowing
calculations to proceed under specific conditions.
- Modified conditions for executing operations on variables, ensuring
compatibility with the framework's operational modes.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: HydrogenSulfate <490868991@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: HydrogenSulfate

deepmd/pd/model/descriptor/se_a.py

@@ -19,6 +19,7 @@
     prod_env_mat,
 )
 from deepmd.pd.utils import (
+    decomp,
     env,
 )
 from deepmd.pd.utils.env import (
@@ -744,7 +745,8 @@ def forward(
                     "Compressed environment is not implemented yet."
                 )
             else:
-                if rr.numel() > 0:
+                # NOTE: control flow with double backward is not supported well yet by paddle.jit
+                if not paddle.framework.in_dynamic_mode() or decomp.numel(rr) > 0:
                     rr = rr * mm.unsqueeze(2).astype(rr.dtype)
                     ss = rr[:, :, :1]
                     if self.compress:

deepmd/pd/model/network/layernorm.py

@@ -100,8 +100,8 @@ def forward(
         yy: paddle.Tensor
             The output.
         """
-        # if xx.numel() > 0:
-        if decomp.numel(xx):
+        # NOTE: control flow with double backward is not supported well yet by paddle.jit
+        if not paddle.framework.in_dynamic_mode() or decomp.numel(xx) > 0:
             variance, mean = (
                 paddle.var(xx, axis=-1, unbiased=False, keepdim=True),
                 paddle.mean(xx, axis=-1, keepdim=True),

deepmd/pd/utils/nlist.py

@@ -7,6 +7,7 @@
 import paddle
 
 from deepmd.pd.utils import (
+    decomp,
     env,
 )
 from deepmd.pd.utils.region import (
@@ -97,7 +98,9 @@ def build_neighbor_list(
     nall = coord.shape[1] // 3
     # fill virtual atoms with large coords so they are not neighbors of any
     # real atom.
-    if coord.numel() > 0:
+
+    # NOTE: control flow with double backward is not supported well yet by paddle.jit
+    if not paddle.framework.in_dynamic_mode() or decomp.numel(coord) > 0:
         xmax = paddle.max(coord) + 2.0 * rcut
     else:
         xmax = paddle.zeros([], dtype=coord.dtype).to(device=coord.place) + 2.0 * rcut
@@ -240,7 +243,8 @@ def build_directional_neighbor_list(
     nall_neig = coord_neig.shape[1] // 3
     # fill virtual atoms with large coords so they are not neighbors of any
     # real atom.
-    if coord_neig.numel() > 0:
+    # NOTE: control flow with double backward is not supported well yet by paddle.jit
+    if not paddle.framework.in_dynamic_mode() or decomp.numel(coord_neig) > 0:
         xmax = paddle.max(coord_cntl) + 2.0 * rcut
     else:
         xmax = (