fix(loss): exclude mixed_type padding atoms from the training loss (#5738)

## Problem

In the `mixed_type` data format, short frames are padded with `type =
-1` ghost atoms up to a fixed `nloc`, and the real atom count varies per
frame within a batch. The training loss normalized by the padded scalar
`natoms` and took unmasked or cross-frame-pooled means, so ghost atoms
diluted the force/atomic denominators and mis-normalized the extensive
energy/virial/property terms. As a result a padded `[3-atom + 5-atom]`
batch did not produce the same loss/gradient as processing the 3-atom
and 5-atom frames separately. Only `mixed_type` batches are affected;
non-mixed training was already exact.

## Fix

The per-atom mask (`atype >= 0`) now reaches the loss under the existing
`model_dict["mask"]` convention: pt_expt already propagates it; the pt
(torch.jit) backend recovers it from `atype` via a new
`TaskLoss._inject_atom_mask` helper called from every pt loss `forward`
(the exported forward drops the model's per-atom mask, so it is
recovered training-side only — the exported artifact is untouched).

Every loss term is then normalized per frame so that a padded batch
equals the grad-accumulation of the individual frames at their real
sizes: per-atom terms (force, atom_ener, atom_pref, atomic dos/tensor,
spin real-force, generalized force) use a per-frame masked mean;
extensive terms (energy, virial, property) divide by the per-frame real
atom count; global already-reduced terms (global dos/tensor) use a plain
mean with the previous atom-count weighting dropped. Every change
reduces exactly to the previous formula when the mask is all-ones, so
non-mixed training is numerically identical (no-op to rounding). Covered
across five shared loss types in both backends:
`deepmd/dpmodel/loss/{ener,ener_spin,dos,tensor,property}.py` (which
serve pt_expt) and
`deepmd/pt/loss/{ener,ener_spin,dos,tensor,property}.py`.

Two additional fixes surfaced during the work: the extensive property
normalization called `xp.sum(mask, -1)` with a positional axis, which
raises `TypeError` under the array_api_compat torch namespace (every
pt_expt extensive-property run) — now `axis=-1`; and `ener_spin`'s MAE
energy and real-force terms were pre-existingly inconsistent with
`ener.py` (they summed over frames without per-atom normalization) —
they are now aligned with `ener.py`, which changes their non-mixed MAE
loss values (a deliberate bug fix, see Known Limitations).

## Test

New `source/tests/common/dpmodel/test_loss_padding.py` and
`source/tests/pt/test_loss_padding.py` assert, for every per-atom and
extensive term of all five loss types in both backends, that a padded
`[3+5]` batch loss equals the mean of the two frames processed
separately, plus an all-ones-mask non-mixed no-op guard per term, and a
torch-tensor path through the dpmodel property loss (which reproduces
the positional-axis crash on the old form). An audit added invariant
coverage for the generalized-force and spin magnetic-force terms and
confirmed they are free of padding artifacts.

## Known limitations

The tf backend loss is unchanged and retains the same mixed_type
behavior (follow-up). The pt-only losses `dens`, `population`, `denoise`
are not covered (follow-up). `ener_spin`'s magnetic-force (`force_mag`)
MAE term uses a sum over frames rather than a mean, so it does not
satisfy the frame-average invariant — this is not a padding artifact
(ghost atoms are correctly excluded via `mask_mag`), but a separate
pre-existing MAE frame-normalization inconsistency, left for a follow-up
decision. The `enable_atom_ener_coeff` path sums ghost atomic energies
before the energy reduction (pre-existing; ghost atom_ener is ~0 by
convention). Existing `mixed_type` trainings will not reproduce
numerically — the new values are the correct ones. Ghost label forces
are assumed ~0 by the dpdata convention; the mask makes the loss robust
even if they are not.

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

* **New Features**
* Added automatic atom-mask injection so padded/multi-frame inputs with
ghost atoms are handled consistently.

* **Bug Fixes**
* Reworked masked loss normalization to use per-frame masked-mean
reductions for energy, forces, virials, atom/property terms, DOS/CDF,
tensor L2, and spin losses.
* Standardized masked global DOS/CDF and global tensor L2 to use
unweighted mean squared error.
* Improved masking behavior for generalized-force projection and updated
RMSE/MAE reporting accordingly.

* **Tests**
* Added/extended gradient-accumulation and padding-mask invariance
suites across dpmodel and pt backends, including atom-mask injection
coverage.
<!-- 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

@@ -130,15 +130,17 @@ def call(
             local_label = xp.reshape(
                 label_dict["atom_dos"], (-1, natoms, self.numb_dos)
             )
-            diff = xp.reshape(local_pred - local_label, (-1, self.numb_dos))
+            diff3d = local_pred - local_label  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                mask = xp.reshape(model_dict["mask"], (-1,))
-                mask_float = xp.astype(mask, diff.dtype)
-                diff = diff * mask_float[:, None]
-                n_valid = xp.sum(mask_float)
-                l2_local_loss_dos = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos)
+                # 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)
             else:
-                l2_local_loss_dos = xp.mean(xp.square(diff))
+                l2_local_loss_dos = xp.mean(xp.square(diff3d))
             loss += pref_ados * l2_local_loss_dos
             more_loss["rmse_local_dos"] = self.display_if_exist(
                 xp.sqrt(l2_local_loss_dos), find_local
@@ -155,15 +157,17 @@ def call(
                 xp.reshape(label_dict["atom_dos"], (-1, natoms, self.numb_dos)),
                 axis=-1,
             )
-            diff = xp.reshape(local_pred_cdf - local_label_cdf, (-1, self.numb_dos))
+            diff3d = local_pred_cdf - local_label_cdf  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                mask = xp.reshape(model_dict["mask"], (-1,))
-                mask_float = xp.astype(mask, diff.dtype)
-                diff = diff * mask_float[:, None]
-                n_valid = xp.sum(mask_float)
-                l2_local_loss_cdf = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos)
+                # 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)
             else:
-                l2_local_loss_cdf = xp.mean(xp.square(diff))
+                l2_local_loss_cdf = xp.mean(xp.square(diff3d))
             loss += pref_acdf * l2_local_loss_cdf
             more_loss["rmse_local_cdf"] = self.display_if_exist(
                 xp.sqrt(l2_local_loss_cdf), find_local
@@ -175,15 +179,14 @@ def call(
             global_pred = xp.reshape(model_dict["dos"], (-1, self.numb_dos))
             global_label = xp.reshape(label_dict["dos"], (-1, self.numb_dos))
             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.sum(model_dict["mask"], axis=-1, keepdims=True)
-                l2_global_loss_dos = xp.mean(
-                    xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num)
+                atom_num = xp.mean(
+                    xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
                 )
-                atom_num = xp.mean(xp.astype(atom_num, diff.dtype))
             else:
                 atom_num = natoms
-                l2_global_loss_dos = xp.mean(xp.square(diff))
             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
@@ -199,15 +202,14 @@ def call(
                 xp.reshape(label_dict["dos"], (-1, self.numb_dos)), axis=-1
             )
             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.sum(model_dict["mask"], axis=-1, keepdims=True)
-                l2_global_loss_cdf = xp.mean(
-                    xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num)
+                atom_num = xp.mean(
+                    xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
                 )
-                atom_num = xp.mean(xp.astype(atom_num, diff.dtype))
             else:
                 atom_num = natoms
-                l2_global_loss_cdf = xp.mean(xp.square(diff))
             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