Merge upstream/master into feat-graph-dpa2

Conflict: source/api_cc/src/common.cc -- both sides fixed the same
daparam-stride aparam buffer sizing (deepmodeling#5790 upstream, e95f2ba here);
code identical, comment text taken from upstream so the hunk drops out
of this PR's diff. deepmd/dpmodel/descriptor/dpa1.py auto-merged with
master's deepmodeling#5785 (dense call is unconditionally _call_dense; verified
post-merge).

Author: Han Wang

.github/workflows/test_cc.yml

@@ -76,7 +76,7 @@ jobs:
         if: matrix.enable_paddle
       - run: |
           export LD_LIBRARY_PATH=${{ github.workspace }}/dp_test/lib:$LD_LIBRARY_PATH
-          pytest --cov=deepmd source/lmp/tests
+          source/install/retry_on_sigterm.sh 3 pytest --cov=deepmd source/lmp/tests
         env:
           OMP_NUM_THREADS: 1
           TF_INTRA_OP_PARALLELISM_THREADS: 1

deepmd/dpmodel/descriptor/dpa1.py

@@ -613,18 +613,25 @@ def call(
         sw
             The smooth switch function.
         """
-        xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist)
-        nloc = nlist.shape[1]
-        nall = xp.reshape(coord_ext, (nlist.shape[0], -1)).shape[1] // 3
-        # graph-eligible configs route through the graph-native adapter (decision
-        # #14: graph = single math source, dense call = thin adapter). Ineligible
-        # configs (compressed descriptors) and the ghost case with no mapping
-        # fall back to the legacy dense body. The graph needs `mapping` to fold
-        # ghosts to local owners; without it only nall == nloc is valid.
-        if self.uses_graph_lower() and (mapping is not None or nall == nloc):
-            return self._call_graph_adapter(coord_ext, atype_ext, nlist, mapping)
-        else:
-            return self._call_dense(coord_ext, atype_ext, nlist)
+        # The dense ``call`` always runs the legacy dense body -- it is the
+        # cross-backend consistency reference and must match the tf/pt/pd/jax
+        # dense descriptors bit-for-bit. It previously routed graph-eligible
+        # configs through ``_call_graph_adapter`` (decision #14), but the
+        # adapter is bit-exact ONLY in the trivial-statistics regime
+        # (``davg == 0``): the dense se_atten body leaks a phantom
+        # padding-neighbor ``-davg/dstd`` residual (``EnvMat.call`` subtracts
+        # ``davg`` AFTER the padding rows' geometry is weight-zeroed, and
+        # neither the empty ``exclude_types`` mask nor the attention layers
+        # re-mask it) that the graph path deliberately omits -- the graph
+        # output is the physically correct one, but ``call`` must reproduce
+        # the dense reference. The former gate also made ``mapping`` -- an
+        # argument that only enables ghost folding on graph routes -- silently
+        # change the numerics of a dense call. The graph-native route is
+        # reached exclusively through :meth:`call_graph` (pt_expt
+        # ``forward_atomic_graph`` and the graph ``.pt2``), never through
+        # ``call``. ``_call_graph_adapter`` is retained as the
+        # bit-exact-regime reference exercised by the adapter parity tests.
+        return self._call_dense(coord_ext, atype_ext, nlist)
 
     def _call_graph_adapter(
         self,
@@ -633,14 +640,27 @@ def _call_graph_adapter(
         nlist: Array,
         mapping: Array | None,
     ) -> Array:
-        """Regime-1 dense->graph adapter (the eligible ``call`` path).
+        """Regime-1 dense->graph adapter.
 
         Builds a NeighborGraph from the dense quartet with the SHAPE-STATIC
         converter (``compact=False``, so this is jit/export-traceable -- no
         ``nonzero``), runs :meth:`call_graph`, and reconstructs the dense-shaped
-        ``sw``. Preserves the dense 5-tuple ABI exactly; masked invalid edges
-        contribute zero in ``call_graph``'s ``segment_sum`` so the output is
-        identical to the legacy dense body.
+        ``sw``. Preserves the dense 5-tuple ABI; masked invalid edges
+        contribute zero in ``call_graph``'s ``segment_sum``.
+
+        Bit-exact vs :meth:`_call_dense` **only in the trivial-statistics
+        regime** (``davg == 0``). For nonzero ``davg`` the dense body leaks a
+        phantom padding-neighbor ``-davg/dstd`` residual into every padding
+        slot (``EnvMat.call`` subtracts ``davg`` AFTER the padding geometry is
+        weight-zeroed; with empty ``exclude_types`` nothing re-masks it, at
+        any ``attn_layer``) that the graph path deliberately omits. The graph
+        kernel additionally applies the slot-0 statistics ``[:, 0, :]`` to
+        every edge -- exact for real stat-computed tables (slot-uniform by
+        construction), not for artificially slot-varying ones. This is why
+        the dense :meth:`call` does NOT route here: it is the cross-backend
+        consistency reference. This method is retained as the
+        bit-exact-regime reference exercised by the adapter parity tests; the
+        production graph route is :meth:`call_graph`.
 
         Parameters
         ----------

deepmd/dpmodel/loss/dos.py

@@ -11,6 +11,10 @@
 from deepmd.dpmodel.loss.loss import (
     Loss,
 )
+from deepmd.dpmodel.loss.reduction import (
+    masked_atom_mean,
+    masked_atom_num,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -134,11 +138,9 @@ def call(
             if "mask" in model_dict:
                 # idiom 1: per-frame masked mean, then average over frames
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
-                nf = diff3d.shape[0]
-                sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1))
-                per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1)  # [nf]
-                per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos  # [nf]
-                l2_local_loss_dos = xp.mean(per_frame_sum / per_frame_dof)
+                l2_local_loss_dos = masked_atom_mean(
+                    xp.square(diff3d), maskf, self.numb_dos
+                )
             else:
                 l2_local_loss_dos = xp.mean(xp.square(diff3d))
             loss += pref_ados * l2_local_loss_dos
@@ -161,11 +163,9 @@ def call(
             if "mask" in model_dict:
                 # idiom 1: per-frame masked mean, then average over frames
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
-                nf = diff3d.shape[0]
-                sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1))
-                per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1)  # [nf]
-                per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos  # [nf]
-                l2_local_loss_cdf = xp.mean(per_frame_sum / per_frame_dof)
+                l2_local_loss_cdf = masked_atom_mean(
+                    xp.square(diff3d), maskf, self.numb_dos
+                )
             else:
                 l2_local_loss_cdf = xp.mean(xp.square(diff3d))
             loss += pref_acdf * l2_local_loss_cdf
@@ -181,12 +181,7 @@ def call(
             diff = global_pred - global_label
             # idiom 3: global dos is already padding-invariant; plain mean suffices
             l2_global_loss_dos = xp.mean(xp.square(diff))
-            if "mask" in model_dict:
-                atom_num = xp.mean(
-                    xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
-                )
-            else:
-                atom_num = natoms
+            atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype)
             loss += pref_dos * l2_global_loss_dos
             more_loss["rmse_global_dos"] = self.display_if_exist(
                 xp.sqrt(l2_global_loss_dos) / atom_num, find_global
@@ -204,12 +199,7 @@ def call(
             diff = global_pred_cdf - global_label_cdf
             # idiom 3: global cdf is already padding-invariant; plain mean suffices
             l2_global_loss_cdf = xp.mean(xp.square(diff))
-            if "mask" in model_dict:
-                atom_num = xp.mean(
-                    xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
-                )
-            else:
-                atom_num = natoms
+            atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype)
             loss += pref_cdf * l2_global_loss_cdf
             more_loss["rmse_global_cdf"] = self.display_if_exist(
                 xp.sqrt(l2_global_loss_cdf) / atom_num, find_global

deepmd/dpmodel/loss/ener.py

@@ -11,6 +11,10 @@
 from deepmd.dpmodel.loss.loss import (
     Loss,
 )
+from deepmd.dpmodel.loss.reduction import (
+    masked_atom_mean,
+    per_frame_component_mean,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -296,7 +300,7 @@ def call(
                 if maskf is not None:
                     # Idiom 2 (extensive): per-frame normalization by real-atom count.
                     se = xp.square(energy - energy_hat)  # [nf, k]
-                    per_frame = xp.mean(xp.reshape(se, (_nf, -1)), axis=-1)  # [nf]
+                    per_frame = per_frame_component_mean(se)  # [nf]
                     if not self.use_huber:
                         loss += pref_e * xp.mean(per_frame * inv**norm_exp)
                     else:
@@ -327,9 +331,7 @@ def call(
                 l1_ener_loss = xp.mean(xp.abs(energy - energy_hat))
                 if maskf is not None:
                     abs_e = xp.abs(energy - energy_hat)  # [nf, k]
-                    per_frame_ae = xp.mean(
-                        xp.reshape(abs_e, (_nf, -1)), axis=-1
-                    )  # [nf]
+                    per_frame_ae = per_frame_component_mean(abs_e)  # [nf]
                     l1_ener_masked = xp.mean(per_frame_ae * inv)
                     loss += pref_e * l1_ener_masked
                     more_loss["mae_e"] = self.display_if_exist(
@@ -346,8 +348,7 @@ def call(
                 )
             if mae:
                 if maskf is not None:
-                    abs_e = xp.abs(energy - energy_hat)
-                    per_frame_ae = xp.mean(xp.reshape(abs_e, (_nf, -1)), axis=-1)
+                    per_frame_ae = per_frame_component_mean(xp.abs(energy - energy_hat))
                     mae_e = xp.mean(per_frame_ae * inv)
                 else:
                     mae_e = xp.mean(xp.abs(energy - energy_hat)) * atom_norm_ener
@@ -362,10 +363,7 @@ def call(
                     diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))  # [nf, nloc, 3]
                     maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))  # [nf, nloc, 1]
                     # Masked MSE computed for rmse_f display regardless of use_huber.
-                    sq_f = xp.square(diff_f_3d) * maskf_col  # [nf, nloc, 3]
-                    _pfs = xp.sum(xp.reshape(sq_f, (_nf, -1)), axis=-1)  # [nf]
-                    _pfd = xp.sum(maskf, axis=-1) * 3  # [nf]
-                    l2_force_masked = xp.mean(_pfs / _pfd)
+                    l2_force_masked = masked_atom_mean(xp.square(diff_f_3d), maskf, 3)
                     if not self.use_huber:
                         loss += pref_f * l2_force_masked
                     else:
@@ -435,12 +433,8 @@ def call(
             elif self.loss_func == "mae":
                 if maskf is not None:
                     diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
-                    maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
                     if not self.f_use_norm:
-                        abs_f = xp.abs(diff_f_3d) * maskf_col  # [nf, nloc, 3]
-                        per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1)
-                        per_frame_dof = xp.sum(maskf, axis=-1) * 3
-                        l1_force_masked = xp.mean(per_frame_sum / per_frame_dof)
+                        l1_force_masked = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3)
                     else:
                         diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3))
                         norm_2d = xp.reshape(
@@ -474,11 +468,7 @@ def call(
             if mae:
                 if maskf is not None:
                     diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
-                    maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
-                    abs_f = xp.abs(diff_f_3d) * maskf_col
-                    per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1)
-                    per_frame_dof = xp.sum(maskf, axis=-1) * 3
-                    mae_f = xp.mean(per_frame_sum / per_frame_dof)
+                    mae_f = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3)
                 else:
                     mae_f = xp.mean(xp.abs(diff_f))
                 more_loss["mae_f"] = self.display_if_exist(mae_f, find_force)
@@ -494,7 +484,7 @@ def call(
                     v2d = xp.reshape(virial, (_nf, 9))
                     v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
                     se_v = xp.square(v_hat_2d - v2d)  # [nf, 9]
-                    per_frame_v = xp.mean(se_v, axis=-1)  # [nf]
+                    per_frame_v = per_frame_component_mean(se_v)  # [nf]
                     if not self.use_huber:
                         loss += pref_v * xp.mean(per_frame_v * inv**norm_exp)
                     else:
@@ -526,8 +516,9 @@ def call(
                 if maskf is not None:
                     v2d = xp.reshape(virial, (_nf, 9))
                     v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
-                    abs_v = xp.abs(v_hat_2d - v2d)  # [nf, 9]
-                    per_frame_v = xp.mean(abs_v, axis=-1)  # [nf]
+                    per_frame_v = per_frame_component_mean(
+                        xp.abs(v_hat_2d - v2d)
+                    )  # [nf]
                     l1_virial_masked = xp.mean(per_frame_v * inv)
                     loss += pref_v * l1_virial_masked
                     more_loss["mae_v"] = self.display_if_exist(
@@ -546,8 +537,7 @@ def call(
                 if maskf is not None:
                     v2d = xp.reshape(virial, (_nf, 9))
                     v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
-                    abs_v = xp.abs(v_hat_2d - v2d)
-                    per_frame_v = xp.mean(abs_v, axis=-1)
+                    per_frame_v = per_frame_component_mean(xp.abs(v_hat_2d - v2d))
                     mae_v = xp.mean(per_frame_v * inv)
                 else:
                     mae_v = (
@@ -565,10 +555,10 @@ def call(
                     # Idiom 1 (per-atom masked mean, ncomp=1).
                     ae_2d = xp.reshape(atom_ener, (_nf, _nloc))
                     ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc))
-                    sq_ae = xp.square(ae_hat_2d - ae_2d) * maskf  # [nf, nloc]
-                    per_frame_sum = xp.sum(sq_ae, axis=-1)  # [nf]
                     per_frame_dof = xp.sum(maskf, axis=-1)  # [nf]
-                    l2_ae_masked = xp.mean(per_frame_sum / per_frame_dof)
+                    l2_ae_masked = masked_atom_mean(
+                        xp.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1
+                    )
                     if not self.use_huber:
                         loss += pref_ae * l2_ae_masked
                     else:
@@ -609,10 +599,9 @@ def call(
                 if maskf is not None:
                     ae_2d = xp.reshape(atom_ener, (_nf, _nloc))
                     ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc))
-                    abs_ae = xp.abs(ae_hat_2d - ae_2d) * maskf  # [nf, nloc]
-                    per_frame_sum = xp.sum(abs_ae, axis=-1)  # [nf]
-                    per_frame_dof = xp.sum(maskf, axis=-1)  # [nf]
-                    l1_ae_masked = xp.mean(per_frame_sum / per_frame_dof)
+                    l1_ae_masked = masked_atom_mean(
+                        xp.abs(ae_hat_2d - ae_2d)[:, :, None], maskf, 1
+                    )
                     loss += pref_ae * l1_ae_masked
                     more_loss["mae_ae"] = self.display_if_exist(
                         l1_ae_masked, find_atom_ener
@@ -637,13 +626,9 @@ def call(
                     # Idiom 1 with pref weight (ncomp=3).
                     diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
                     pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3))
-                    maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
-                    sq_pf = xp.square(diff_f_3d) * pf_3d * maskf_col  # [nf, nloc, 3]
-                    per_frame_sum = xp.sum(
-                        xp.reshape(sq_pf, (_nf, -1)), axis=-1
-                    )  # [nf]
-                    per_frame_dof = xp.sum(maskf, axis=-1) * 3  # [nf]
-                    l2_pf_masked = xp.mean(per_frame_sum / per_frame_dof)
+                    l2_pf_masked = masked_atom_mean(
+                        xp.square(diff_f_3d) * pf_3d, maskf, 3
+                    )
                     loss += pref_pf * l2_pf_masked
                     more_loss["rmse_pf"] = self.display_if_exist(
                         xp.sqrt(l2_pf_masked), find_atom_pref
@@ -660,11 +645,7 @@ def call(
                 if maskf is not None:
                     diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
                     pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3))
-                    maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
-                    abs_pf = xp.abs(diff_f_3d) * pf_3d * maskf_col  # [nf, nloc, 3]
-                    per_frame_sum = xp.sum(xp.reshape(abs_pf, (_nf, -1)), axis=-1)
-                    per_frame_dof = xp.sum(maskf, axis=-1) * 3
-                    l1_pf_masked = xp.mean(per_frame_sum / per_frame_dof)
+                    l1_pf_masked = masked_atom_mean(xp.abs(diff_f_3d) * pf_3d, maskf, 3)
                     loss += pref_pf * l1_pf_masked
                     more_loss["mae_pf"] = self.display_if_exist(
                         l1_pf_masked, find_atom_pref