Skip to content

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

Open
wanghan-iapcm wants to merge 1 commit intodeepmodeling:masterfrom
wanghan-iapcm:chore-dp-loss-ut
Open

fix(dpmodel): fix natoms[0] bug, einsum, and return type in EnergyLoss#5325
wanghan-iapcm wants to merge 1 commit intodeepmodeling:masterfrom
wanghan-iapcm:chore-dp-loss-ut

Conversation

@wanghan-iapcm
Copy link
Collaborator

@wanghan-iapcm wanghan-iapcm commented Mar 18, 2026

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

- 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)
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Mar 18, 2026

📝 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

Possibly related PRs

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

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

Copy link
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +89 to +91
model_dict, label_dict, natoms = self._make_data()
loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
self.assertIsNotNone(loss)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Use more_loss here or discard it explicitly.

This unpack is currently unused and Ruff reports RUF059. A lightweight assertion keeps the tuple-contract coverage and satisfies the linter.

♻️ Suggested change
         model_dict, label_dict, natoms = self._make_data()
         loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
         self.assertIsNotNone(loss)
+        self.assertIsInstance(more_loss, dict)
As per coding guidelines `**/*.py`: Always run `ruff check .` and `ruff format .` before committing changes or CI will fail.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
model_dict, label_dict, natoms = self._make_data()
loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
self.assertIsNotNone(loss)
model_dict, label_dict, natoms = self._make_data()
loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
self.assertIsNotNone(loss)
self.assertIsInstance(more_loss, dict)
🧰 Tools
🪛 Ruff (0.15.6)

[warning] 90-90: Unpacked variable more_loss is never used

Prefix it with an underscore or any other dummy variable pattern

(RUF059)

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@source/tests/common/dpmodel/test_loss_ener.py` around lines 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.

Comment on lines +143 to +145
model_dict, label_dict, natoms = self._make_data()
loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
self.assertIsNotNone(loss)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Use more_loss here or discard it explicitly.

This unpack is currently unused and Ruff reports RUF059. A small type assertion is enough if you want to keep validating the new return shape in this test too.

♻️ Suggested change
         model_dict, label_dict, natoms = self._make_data()
         loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
         self.assertIsNotNone(loss)
+        self.assertIsInstance(more_loss, dict)
As per coding guidelines `**/*.py`: Always run `ruff check .` and `ruff format .` before committing changes or CI will fail.
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
model_dict, label_dict, natoms = self._make_data()
loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
self.assertIsNotNone(loss)
model_dict, label_dict, natoms = self._make_data()
loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
self.assertIsNotNone(loss)
self.assertIsInstance(more_loss, dict)
🧰 Tools
🪛 Ruff (0.15.6)

[warning] 144-144: Unpacked variable more_loss is never used

Prefix it with an underscore or any other dummy variable pattern

(RUF059)

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@source/tests/common/dpmodel/test_loss_ener.py` around lines 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.

Comment on lines +373 to +379
loss, more_loss = obj.build(
self.learning_rate,
[self.natoms],
predict,
label,
suffix=suffix,
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Discard the unused TF more_loss binding.

build_tf() drops the auxiliary losses immediately, so this unpack just triggers RUF059. Rename it to _ or _more_loss.

♻️ Suggested change
-        loss, more_loss = obj.build(
+        loss, _ = obj.build(
             self.learning_rate,
             [self.natoms],
             predict,
             label,
             suffix=suffix,
As per coding guidelines `**/*.py`: Always run `ruff check .` and `ruff format .` before committing changes or CI will fail.
🧰 Tools
🪛 Ruff (0.15.6)

[warning] 373-373: Unpacked variable more_loss is never used

Prefix it with an underscore or any other dummy variable pattern

(RUF059)

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@source/tests/consistent/loss/test_ener.py` around lines 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.

@dosubot dosubot bot added the bug label Mar 18, 2026
@codecov
Copy link

codecov bot commented Mar 18, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.31%. Comparing base (24e54bf) to head (fb4c1b4).
⚠️ Report is 15 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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants