refactor(loss): extract masked per-frame reduction idioms to cut nested branching (#5783)

Resolves #5768. Follow-up to #5738.

## What this does

The mixed_type padding fix (#5738) added a masked (`maskf is not None`)
branch to every term of the shared losses. Combined with the existing
`has_*` / `mse`|`mae` / `use_huber` conditions, that produced deeply
nested, duplicated per-frame arithmetic repeated across energy, force,
virial, atom_ener, atom_pref, dos, tensor, and again across `ener_spin`,
in both the dpmodel and pt backends.

This PR extracts the three recurring masked reduction "idioms" into a
single shared `array_api_compat` module,
`deepmd/dpmodel/loss/reduction.py`:

- `masked_atom_mean(elem, maskf, ncomp)` — idiom 1: per-atom masked mean
over `ncomp` components (force, atom_ener, atom_pref, local dos/tensor,
spin real-force).
- `per_frame_component_mean(err)` — idiom 2 primitive: per-frame mean
over components; the caller applies the extensive
`1/real_natoms**norm_exp` weighting (energy, virial).
- `masked_atom_num(mask, natoms, dtype)` — idiom 3 companion: the
display-only atom-count divisor for already-reduced global quantities
(global dos/tensor).

Both backends import the same module; the pt backend calls it with torch
tensors, and `array_api_compat`'s torch-namespace dispatch preserves
autograd and produces bit-identical results.

## Numerics are preserved

This is a pure readability/maintainability refactor with no intended
behavior change. Each helper implements only the masked branch; every
call site keeps its original non-masked `else` expression verbatim, so
the "bit-identical for non-mixed batches" guarantee from #5738 holds
(defaulting the mask to all-ones would change the reduction order at the
ULP level and is deliberately avoided). Operand order and the
mask-application point are preserved in every conversion.

## Testing

The safety net is the grad-accumulation-invariant and
no-op-for-non-mixed tests added in #5738 plus the cross-backend
consistency loss tests, all of which stay green:

- `source/tests/common/dpmodel/test_loss_padding.py`,
`source/tests/pt/test_loss_padding.py`
- `source/tests/consistent/loss` (dpmodel <-> pt <-> tf): 244 passed,
176 skipped — unchanged, the load-bearing bit-identity guard
- new direct unit tests for the three helpers (numpy + torch, autograd,
bit-identity vs the torch-native inline form):
`source/tests/common/dpmodel/test_loss_reduction.py`

## Known limitations

- The shared module lives under `deepmd/dpmodel/loss` and is imported by
`deepmd/pt/loss` — the first such cross-backend import in the loss
layer. This is intentional (single source of truth for the reduction
math, consistent with `deepmd/dpmodel` being the designated shared-math
layer), but it does couple the pt loss backend to `array_api_compat`.
- The Huber-masked sub-branches, the `f_use_norm` force variant, the
generalized-force (`gf`) input masking, and the spin magnetic-force
(`mask_mag`) block are intentionally NOT routed through the helpers —
they are not one of the three idioms.
- The TensorFlow (`deepmd/tf`) loss backend is out of scope; it does not
use the mixed_type padding mask.

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

* **Improvements**
* Improved masked loss reductions across energy, forces, virial, tensor,
DOS, and spin training by switching to shared masked/per-frame reduction
helpers.
* Standardized masked atom normalization and per-frame aggregation for
consistent RMSE/MAE reporting in local and global DOS/tensor
computations.
* Preserved behavior and control flow across NumPy/JAX and PyTorch
execution paths.

* **Tests**
* Added coverage for masked mean, per-frame component mean, and
atom-count normalization (including `None` masks), plus PyTorch autograd
checks and no-NaN behavior on all-padding frames.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: wanghan-iapcm

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