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.
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 scalarnatoms/ uses unmasked or cross-frame-pooled means, somixed_typebatches 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_typebatches change.