fix(dpmodel): align Loss.call return type with implementation

[njzjz-bot]

Problem

• deepmd/dpmodel/loss/loss.py declares Loss.call(...) -> dict[str, Array], but the concrete implementation EnergyLoss.call(...) returns (loss, more_loss).
• In EnergyLoss.call generalized-force branch, natoms is annotated as int but code uses natoms[0], which looks like a bug.

Change

• Update abstract Loss.call return type to tuple[Array, dict[str, Array]].
• Update EnergyLoss.call return type hint accordingly.
• Replace natoms[0] with natoms when reshaping force/drdq in generalized-force loss.
• Expand the base-class docstring to document loss vs more_loss.

Notes

• This is API typing/docs alignment + a small bugfix; no intended behavior change besides fixing the incorrect natoms[0] indexing.
• This matches existing usage in pt_expt where loss, more_loss = ....

Authored by OpenClaw (model: gpt-5.2)

[coderabbitai]

📝 Walkthrough

Walkthrough

Updates the Loss API: Loss.call() and EnergyLoss.call() now return a tuple (loss, more_loss) and the EnergyLoss implementation adjusts tensor reshaping (uses natoms * 3 for force/drdq flattening) affecting downstream einsum/metric computations.

Changes

Cohort / File(s)	Summary
Loss API & docs deepmd/dpmodel/loss/loss.py	Changed abstract Loss.call() return type to tuple[Array, dict[str, Array]] and updated docstring to describe scalar loss and more_loss dict.
Energy loss implementation deepmd/dpmodel/loss/ener.py	Updated EnergyLoss.call() signature to return (loss, more_loss); changed internal tensor reshaping to use natoms * 3 (flattening force/drdq) and adjusted subsequent einsum and RMSE-related computations to match new shapes.

Sequence Diagram(s)

(omitted — conditions for diagrams not met)

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related issues

• dpmodel: loss API should consistently return (loss, more_loss) tuple; current signature/docs mismatch #5313: Matches API change to return (loss, more_loss) and the tensor-shape adjustments in EnergyLoss.

Possibly related PRs

• fix(dpmodel): rmse_e loss normalized by n_atoms #5124: Also modifies EnergyLoss.call; related by overlapping changes to return/value shaping.
• fix(dpmodel): fix rmse_* in more_loss #5105: Related changes to RMSE/more_loss outputs and normalization logic.
• fix(dpmodel): fix energy loss #4765: Alters EnergyLoss.call outputs/keys; likely touches the same API surface.

Suggested labels

Docs

Suggested reviewers

• njzjz
• wanghan-iapcm

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: aligning the Loss.call return type annotation with its actual implementation (tuple instead of dict).
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

📝 Coding Plan

• Generate coding plan for human review comments

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[njzjz-bot]

Follow-up: this PR also fixes a likely bug in deepmd/dpmodel/loss/ener.py generalized-force branch where natoms is annotated as int but was indexed as natoms[0]. The code now uses natoms directly when reshaping force/drdq.

— OpenClaw 2026.3.8 (model: gpt-5.2)

[codecov]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.30%. Comparing base (09345bf) to head (5979d62).

Files with missing lines	Patch %	Lines
deepmd/dpmodel/loss/ener.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #5314   +/-   ##
=======================================
  Coverage   82.29%   82.30%           
=======================================
  Files         775      775           
  Lines       77627    77628    +1     
  Branches     3676     3675    -1     
=======================================
+ Hits        63887    63888    +1     
- Misses      12566    12568    +2     
+ Partials     1174     1172    -2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[wanghan-iapcm]

I will open a new PR, fixing the natoms issue and add UT for the dpmodel's loss.

[wanghan-iapcm]

can be replaced by #5325