TF backend loss: exclude mixed_type padding atoms (per-frame normalization)
bug
Follow-up to #5738.
PR #5738 fixed the mixed_type padding-atom dilution in the training loss for the dpmodel / pt / pt_expt backends: every loss term is now normalized per frame so a padded `[3+5]`-atom batch yields the same loss/gradient as processing each frame separately and averaging. The **TensorFlow backend loss was left unchanged** and still normalizes by the padded scalar `natoms` / uses unmasked or cross-frame-pooled means, so `mixed_type` batches remain mis-normalized there.
Scope: apply the same per-frame masked normalization (per-atom masked mean for extensive/atomic terms; per-frame real-atom count for extensive energy/virial/property) to the TF backend loss, with an all-ones-mask no-op guard so non-mixed training stays bit-identical. Add grad-accumulation-invariant tests mirroring `source/tests/*/test_loss_padding.py`.
Non-mixed TF training is unaffected; only `mixed_type` batches change.
0 条评论