fix(dpmodel): fix natoms[0] bug, einsum, and return type in EnergyLoss

[wanghan-iapcm]

• Fix natoms[0] -> natoms in generalized force branch (natoms is int)
• Replace xp.einsum with array-API-compatible xp.sum + broadcasting
• Fix return type annotation of Loss.call and EnergyLoss.call from dict[str, Array] to tuple[Array, dict[str, Array]]
• Add TestEnerGF consistency test for generalized force code path
• Add dpmodel-level unit tests for EnergyLoss (basic, aecoeff, generalized force, huber, serialize round-trip)

Summary by CodeRabbit

• Bug Fixes

  • Enhanced numerical accuracy in energy loss force calculations through optimized computation methods.

• Tests

  • Added comprehensive test coverage for energy loss calculations, including generalized coordinate scenarios.
  • Expanded multi-backend compatibility validation across TensorFlow, PyTorch, JAX, Array API, and Paddle.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR updates the Loss and EnergyLoss class signatures to return a tuple of (loss, auxiliary losses dictionary) instead of a single dictionary, and fixes array indexing and reshaping operations in generalized-force computations to align with array-api-compat semantics. Comprehensive tests validate the updated behavior across multiple backends.

Changes

Cohort / File(s)	Summary
Loss Module deepmd/dpmodel/loss/loss.py, deepmd/dpmodel/loss/ener.py	Updated Loss.call and EnergyLoss.call return types from dict[str, Array] to tuple[Array, dict[str, Array]]. In EnergyLoss, replaced einsum operations with element-wise multiplication and summation, fixed natoms indexing from natoms[0] to natoms, and adjusted force/reshaping operations to 1D per-frame then per-atom/chunk shapes for array-api compatibility.
Unit Tests source/tests/common/dpmodel/test_loss_ener.py	New test module validating EnergyLoss forward passes with standard and atom-energy coefficients, generalized-force paths, Huber loss variants, and serialization round-trips using deterministic seeding.
Consistent Tests source/tests/consistent/loss/test_ener.py	New TestEnerGF test class covering energy loss with generalized coordinates (numb_generalized_coord > 0) across TensorFlow, PyTorch, JAX, Array API Strict, and Paddle backends with setUp, build, eval, and extraction methods.

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: PR directly addresses the API type-signature changes (Loss.call return type) and natoms-indexing inconsistencies described in the issue.

Possibly related PRs

• fix(dpmodel): fix rmse_* in more_loss #5105: Modifies EnergyLoss.call behavior and loss metadata (rmse_\* / more_loss) handling alongside return shape changes.
• fix(pt): Huber energy loss uses raw model energy #5168: Fixes array-shape, einsum, and reshaping logic in energy-loss computation with atom_energy/atom_ener_coeff handling.
• feat: add huber loss #4684: Modifies EnergyLoss.call return shape and force-reshaping logic with Huber-loss support additions.

Suggested labels

bug, Python

Suggested reviewers

• njzjz
• iProzd

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.05% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The PR title accurately describes the three main changes: fixing the natoms[0] bug, replacing einsum with array-API-compatible operations, and updating the return type in EnergyLoss.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

📝 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.}

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@source/tests/common/dpmodel/test_loss_ener.py`:
- Around line 89-91: The test currently unpacks (loss, more_loss) from
loss_fn.call but never uses more_loss which triggers RUF059; update the test in
test_loss_ener.py to either assert something lightweight about more_loss (e.g.,
assertIsNotNone(more_loss) or assertGreaterEqual(len(more_loss), 0)) or
explicitly discard it using an underscore (e.g., _, = ...) so the tuple contract
is honored; locate the unpack at the call to loss_fn.call and make the change to
reference or discard the more_loss variable accordingly.
- Around line 143-145: The test currently unpacks loss, more_loss =
loss_fn.call(1.0, natoms, model_dict, label_dict) but never uses more_loss
(RUF059); fix by either explicitly discarding it (e.g., replace with loss, _ =
loss_fn.call(...)) or add a minimal assertion/type check to use it (e.g.,
self.assertIsNotNone(more_loss) or self.assertIsInstance(more_loss, dict)) so
the return shape is validated; modify the call site in the test where
loss_fn.call is invoked to implement one of these options.

In `@source/tests/consistent/loss/test_ener.py`:
- Around line 373-379: Unpack only the used loss from the obj.build call by
discarding the unused TensorFlow auxiliary binding: change the second variable
name from more_loss to _ or _more_loss in the assignment where obj.build(...) is
called (the invocation of obj.build that currently assigns loss, more_loss) so
the unused TF loss is ignored and RUF059 is avoided; run ruff check . and ruff
format . after the change.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 14d402a1-8899-4080-b0ca-971e9789866b

📥 Commits

Reviewing files that changed from the base of the PR and between b2805fb and fb4c1b4.

📒 Files selected for processing (4)

• deepmd/dpmodel/loss/ener.py
• deepmd/dpmodel/loss/loss.py
• source/tests/common/dpmodel/test_loss_ener.py
• source/tests/consistent/loss/test_ener.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.31%. Comparing base (24e54bf) to head (fb4c1b4).
⚠️ Report is 17 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5325      +/-   ##
==========================================
- Coverage   82.32%   82.31%   -0.01%     
==========================================
  Files         768      775       +7     
  Lines       77098    77627     +529     
  Branches     3659     3675      +16     
==========================================
+ Hits        63469    63899     +430     
- Misses      12458    12554      +96     
- Partials     1171     1174       +3     

☔ 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.