refactor(loss): extract masked per-frame reduction idioms to cut nested branching#5783
Open
wanghan-iapcm wants to merge 9 commits into
Open
refactor(loss): extract masked per-frame reduction idioms to cut nested branching#5783wanghan-iapcm wants to merge 9 commits into
wanghan-iapcm wants to merge 9 commits into
Conversation
added 9 commits
July 13, 2026 13:08
Extract the three per-frame reduction idioms introduced by the mixed_type padding fix (deepmodeling#5738) into deepmd/dpmodel/loss/reduction.py, written with array_api_compat so both the dpmodel and pt loss backends can share them. Callers still keep the original non-masked expression in the else branch, so the bit-identical-for-non-mixed guarantee is preserved. No call sites adopt the helpers yet. Ref deepmodeling#5768
Collapse the idiom-1 (force/atom-energy/prefactor-force) and idiom-2 (energy) masked per-frame arithmetic in deepmd/dpmodel/loss/ener.py to calls into the shared reduction helpers. Huber, f_use_norm, generalized force and the non-masked else branches are untouched; numerics are bit-identical. Ref deepmodeling#5768
Complete the dpmodel EnergyLoss adoption by routing the virial masked per-frame branches through per_frame_component_mean, matching the energy term and EnergySpinLoss. Numerics bit-identical; else/huber/display branches untouched. Ref deepmodeling#5768
EnergySpinLoss now reuses the same reduction helpers as EnergyLoss for its energy/real-force/atom-energy/virial masked branches, cutting the largest duplication between the two files. The magnetic-force (mask_mag) block is spin-specific and left bespoke. Numerics bit-identical. Ref deepmodeling#5768
…loss Local dos/cdf/tensor masked branches use the idiom-1 helper; the global dos/cdf/tensor terms keep their plain-mean loss and only route the display-only atom-count divisor through masked_atom_num. Numerics bit-identical. Ref deepmodeling#5768
…Loss deepmd/pt/loss/ener.py now imports the array_api_compat reduction helpers from deepmd.dpmodel.loss.reduction and calls them with torch tensors; the torch namespace dispatch preserves autograd and is bit-identical to the prior inlined torch code. Huber / f_use_norm / generalized-force paths and the .to(GLOBAL_PT_FLOAT_PRECISION) casts are unchanged. Ref deepmodeling#5768
…ng, test) Address final-review nits on the shared reduction module: per_frame_component_mean now derives nf from err.shape[0] internally (matching masked_atom_mean and removing a caller footgun), add the missing Returns type line to masked_atom_num's docstring, and add an autograd test for per_frame_component_mean. Numerics bit-identical. Ref deepmodeling#5768
Contributor
📝 WalkthroughWalkthroughThe change adds shared masked-reduction helpers, tests their NumPy and PyTorch behavior, and applies them across dpmodel and PyTorch DOS, energy, spin-energy, and tensor losses. ChangesMasked loss reduction refactor
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_loss_reduction.py (1)
19-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_refduplicates the implementation formula rather than independently verifying it.
_refmirrorsmasked_atom_mean's exact computation, sotest_numpy_matches_referencemainly confirms the numpy dispatch path matches itself rather than validating correctness against an independently derived expected value. Consider adding at least one case with a manually computed expected scalar to catch a shared logic bug in both.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/common/dpmodel/test_loss_reduction.py` around lines 19 - 33, Replace or supplement the formula-based _ref oracle in test_numpy_matches_reference with at least one test case whose expected reduction scalar is manually computed and asserted directly. Keep coverage for both ncomp values as appropriate, but ensure the test can detect a shared error in masked_atom_mean and _ref rather than comparing two identical implementations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/loss/reduction.py`:
- Around line 30-55: Update masked_atom_mean to guard zero values in
per_frame_dof before dividing per_frame_sum by it, assigning a finite neutral
per-frame result for all-padding frames while preserving the existing
calculation for positive degrees of freedom and the final frame mean.
---
Nitpick comments:
In `@source/tests/common/dpmodel/test_loss_reduction.py`:
- Around line 19-33: Replace or supplement the formula-based _ref oracle in
test_numpy_matches_reference with at least one test case whose expected
reduction scalar is manually computed and asserted directly. Keep coverage for
both ncomp values as appropriate, but ensure the test can detect a shared error
in masked_atom_mean and _ref rather than comparing two identical
implementations.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 611b6d93-b490-4ef1-b0e6-52b9b72acb1c
📒 Files selected for processing (10)
deepmd/dpmodel/loss/dos.pydeepmd/dpmodel/loss/ener.pydeepmd/dpmodel/loss/ener_spin.pydeepmd/dpmodel/loss/reduction.pydeepmd/dpmodel/loss/tensor.pydeepmd/pt/loss/dos.pydeepmd/pt/loss/ener.pydeepmd/pt/loss/ener_spin.pydeepmd/pt/loss/tensor.pysource/tests/common/dpmodel/test_loss_reduction.py
Comment on lines
+30
to
+55
| def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array: | ||
| """Idiom 1: per-atom masked mean over ``ncomp`` components, averaged over frames. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| elem : Array | ||
| Non-negative per-element contribution of shape ``[nf, nloc, ncomp]`` | ||
| (already squared or abs'd, and pre-multiplied by any per-atom weight | ||
| such as ``atom_pref``). NOT yet multiplied by the mask. | ||
| maskf : Array | ||
| Per-atom real/ghost mask of shape ``[nf, nloc]`` (1.0 real, 0.0 ghost). | ||
| ncomp : int | ||
| Number of components per atom (force: 3, atom energy: 1, | ||
| dos: ``numb_dos``, tensor: ``tensor_size``). | ||
|
|
||
| Returns | ||
| ------- | ||
| Array | ||
| ``mean_over_frames( sum(elem * mask) / (real_natoms * ncomp) )``. | ||
| """ | ||
| xp = array_api_compat.array_namespace(elem, maskf) | ||
| nf = elem.shape[0] | ||
| masked = elem * maskf[:, :, None] | ||
| per_frame_sum = xp.sum(xp.reshape(masked, (nf, -1)), axis=-1) | ||
| per_frame_dof = xp.sum(maskf, axis=-1) * ncomp | ||
| return xp.mean(per_frame_sum / per_frame_dof) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
set -euo pipefail
printf '\n== reduction.py outline ==\n'
ast-grep outline deepmd/dpmodel/loss/reduction.py --view expanded || true
printf '\n== masked_atom_mean call sites ==\n'
rg -n "masked_atom_mean\(" deepmd -S || true
printf '\n== relevant mask logic and frame-count invariants ==\n'
rg -n "real atoms|real_natoms|maskf|atom mask|ghost|padding|nloc|nframes|all-padding|zero real atoms" deepmd -S || trueRepository: deepmodeling/deepmd-kit
Length of output: 50380
🏁 Script executed:
set -euo pipefail
printf '\n== reduction.py excerpt ==\n'
sed -n '1,140p' deepmd/dpmodel/loss/reduction.py
printf '\n== call sites with context ==\n'
rg -n -C 3 "masked_atom_mean\(" deepmd -S
printf '\n== mask generation / invariants candidates ==\n'
rg -n -C 3 "maskf|real_natoms|natoms|ghost|padding" deepmd/dpmodel deepmd/utils deepmd/pt -SRepository: deepmodeling/deepmd-kit
Length of output: 50380
🏁 Script executed:
set -euo pipefail
printf '\n== loss/ener.py mask handling excerpt ==\n'
sed -n '220,430p' deepmd/dpmodel/loss/ener.py
printf '\n== dpmodel mask producers / consumers ==\n'
rg -n -C 2 '"mask"|maskf|mask_mag|find_atom_|real_natoms' deepmd/dpmodel deepmd/pt/model -S | sed -n '1,220p'Repository: deepmodeling/deepmd-kit
Length of output: 26677
🏁 Script executed:
python3 - <<'PY'
import math
import numpy as np
try:
import torch
except Exception:
torch = None
# Numpy behavior for 0/0 in the helper shape
elem = np.zeros((1, 2, 3), dtype=np.float32)
mask = np.array([[0.0, 0.0]], dtype=np.float32)
masked = elem * mask[:, :, None]
per_frame_sum = masked.reshape(1, -1).sum(axis=-1)
per_frame_dof = mask.sum(axis=-1) * 3
print("numpy per_frame_sum:", per_frame_sum)
print("numpy per_frame_dof:", per_frame_dof)
print("numpy ratio:", per_frame_sum / per_frame_dof)
if torch is not None:
elem_t = torch.zeros((1, 2, 3), dtype=torch.float32)
mask_t = torch.zeros((1, 2), dtype=torch.float32)
masked_t = elem_t * mask_t[:, :, None]
per_frame_sum_t = masked_t.reshape(1, -1).sum(dim=-1)
per_frame_dof_t = mask_t.sum(dim=-1) * 3
print("torch per_frame_sum:", per_frame_sum_t)
print("torch per_frame_dof:", per_frame_dof_t)
print("torch ratio:", per_frame_sum_t / per_frame_dof_t)
PY
rg -n -C 2 "at least one real atom|real_natoms|sum\\(maskf|mask\\s*=|0 real atoms|zero real atoms|non-empty|empty frame" deepmd/dpmodel deepmd/pt -S | sed -n '1,220p'Repository: deepmodeling/deepmd-kit
Length of output: 16293
Guard zero-real-atom frames in masked_atom_mean.
per_frame_dof becomes 0 for an all-padding frame, so sum / dof turns into NaN and can poison every loss that uses this helper. A zero-dof guard here would keep normal behavior unchanged.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@deepmd/dpmodel/loss/reduction.py` around lines 30 - 55, Update
masked_atom_mean to guard zero values in per_frame_dof before dividing
per_frame_sum by it, assigning a finite neutral per-frame result for all-padding
frames while preserving the existing calculation for positive degrees of freedom
and the final frame mean.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5783 +/- ##
==========================================
- Coverage 79.77% 79.61% -0.17%
==========================================
Files 1020 1021 +1
Lines 116872 116763 -109
Branches 4308 4303 -5
==========================================
- Hits 93239 92956 -283
- Misses 22095 22265 +170
- Partials 1538 1542 +4 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
approved these changes
Jul 13, 2026
njzjz-bot
left a comment
njzjz-bot
left a comment
Contributor
There was a problem hiding this comment.
Reviewed commit 13b74c9. I found no blocking correctness issues in the shared reduction helpers or their dpmodel and PyTorch call sites. The refactor preserves the existing shape, dtype, masking, normalization, and reduction order.
Local validation:
- common helper and padding tests: 65 passed
- PyTorch padding tests: 58 passed
- cross-backend loss consistency tests: 232 passed, 188 skipped
- Ruff check and format check: passed
Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
njzjz
approved these changes
Jul 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 existinghas_*/mse|mae/use_huberconditions, that produced deeply nested, duplicated per-frame arithmetic repeated across energy, force, virial, atom_ener, atom_pref, dos, tensor, and again acrossener_spin, in both the dpmodel and pt backends.This PR extracts the three recurring masked reduction "idioms" into a single shared
array_api_compatmodule,deepmd/dpmodel/loss/reduction.py:masked_atom_mean(elem, maskf, ncomp)— idiom 1: per-atom masked mean overncompcomponents (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 extensive1/real_natoms**norm_expweighting (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
elseexpression 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.pysource/tests/consistent/loss(dpmodel <-> pt <-> tf): 244 passed, 176 skipped — unchanged, the load-bearing bit-identity guardsource/tests/common/dpmodel/test_loss_reduction.pyKnown limitations
deepmd/dpmodel/lossand is imported bydeepmd/pt/loss— the first such cross-backend import in the loss layer. This is intentional (single source of truth for the reduction math, consistent withdeepmd/dpmodelbeing the designated shared-math layer), but it does couple the pt loss backend toarray_api_compat.f_use_normforce 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.deepmd/tf) loss backend is out of scope; it does not use the mixed_type padding mask.Summary by CodeRabbit
Improvements
Tests