fix(dpmodel): keep output reductions on active backend

[njzjz]

Summary

• translate the configured NumPy energy precision into the dtype object of the active array namespace;
• use namespace-native casts and reductions in the dense dpmodel output transform;
• apply the same dtype correction to the ragged graph transform, which otherwise passes a NumPy dtype class to Torch;
• add focused Torch regressions for extensive, masked-intensive, and graph reductions.

Why existing tests missed this

There were no direct tests exercising the backend-agnostic dpmodel transform helpers with Torch arrays. Normal PyTorch and Paddle model tests call their backend-specific transform implementations instead, while the existing dpmodel graph tests use NumPy. As a result, neither Tensor.astype nor the NumPy dtype passed into the Torch namespace was reached.

The new tests invoke the generic dpmodel helpers directly and verify that results remain Torch tensors in the configured energy precision.

Validation

• new transform regressions plus existing ragged graph tests: 10 passed;
• additional reviewer probes passed for NumPy, Torch, JAX, and array-api-strict;
• ruff format .;
• ruff check .;
• git diff --check.

Closes #5640.

Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh

Summary by CodeRabbit

• Bug Fixes

  • Improved energy reduction accuracy by using backend-specific precision for accumulation and per-frame/per-node division, instead of a fixed global dtype.
  • Ensured intensive/extensive reductions consistently cast reducible values—and any provided masks or computed divisors—to the correct backend dtype before summing/averaging.

• Tests

  • Added Torch backend-preservation tests for both dense and graph reduction paths, covering intensive vs extensive behavior with and without masks and validating dtype promotion and correctness.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro Plus

Run ID: a1759cd7-9b43-41c4-b3d9-a4912b2c42b0

📥 Commits

Reviewing files that changed from the base of the PR and between d3ed641 and e900967.

📒 Files selected for processing (1)

• deepmd/dpmodel/model/edge_transform_output.py

🚧 Files skipped from review as they are similar to previous changes (1)

• deepmd/dpmodel/model/edge_transform_output.py

---

📝 Walkthrough

Walkthrough

The output transforms resolve energy precision through the active array namespace before reducing values and masks. New Torch tests cover dense and graph extensive and intensive reductions, including masked and count-based denominator paths.

Changes

Backend-aware output reductions

Layer / File(s)	Summary
Dense backend-aware reductions deepmd/dpmodel/model/transform_output.py, source/tests/common/dpmodel/test_transform_output.py	Dense reductions use namespace-specific energy dtypes for values and masks, with Torch tests covering extensive and intensive outputs with and without masks.
Graph backend-aware reductions deepmd/dpmodel/model/edge_transform_output.py, source/tests/common/dpmodel/test_transform_output.py	Graph reductions cast values and intensive denominators through the active namespace, with Torch coverage for masked, unmasked, extensive, and intensive paths.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: copilot, iprozd, outisli

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title concisely reflects the main change: keeping output reductions on the active backend.
Linked Issues check	✅ Passed	The dense reduction fixes and Torch regression tests directly address the backend-API cast and reduction issues in #5640.
Out of Scope Changes check	✅ Passed	The graph-transform dtype updates still align with the same backend-preservation goal and do not introduce unrelated scope.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

source/tests/common/dpmodel/test_transform_output.py (1)

35-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Gracefully skip tests if Torch is not installed.

If torch is an optional dependency for local testing, consider using pytest.importorskip to gracefully skip the test instead of failing with a ModuleNotFoundError.

💡 Proposed refactor

Apply this change to the current test and the other tests in this file:

-    import torch
+    torch = pytest.importorskip("torch")

Also, remember to add import pytest at the top of the file:

import pytest

🤖 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_transform_output.py` around lines 35 - 37,
Update the Torch-dependent tests in this file, including
test_dense_torch_extensive_reduction_stays_on_backend, to use
pytest.importorskip for the torch import so they are skipped when Torch is
unavailable; add the pytest import at module scope and preserve the existing
test behavior when Torch is installed.

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

Nitpick comments:
In `@source/tests/common/dpmodel/test_transform_output.py`:
- Around line 35-37: Update the Torch-dependent tests in this file, including
test_dense_torch_extensive_reduction_stays_on_backend, to use
pytest.importorskip for the torch import so they are skipped when Torch is
unavailable; add the pytest import at module scope and preserve the existing
test behavior when Torch is installed.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 7cb7c430-ddcb-4944-b3c9-d1b50ff26b2c

📥 Commits

Reviewing files that changed from the base of the PR and between 7d5ad38 and 764d4ca.

📒 Files selected for processing (3)

• deepmd/dpmodel/model/edge_transform_output.py
• deepmd/dpmodel/model/transform_output.py
• source/tests/common/dpmodel/test_transform_output.py

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.63%. Comparing base (5e42987) to head (e900967).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5789      +/-   ##
==========================================
- Coverage   78.88%   78.63%   -0.25%     
==========================================
  Files        1054     1054              
  Lines      121756   121757       +1     
  Branches     4402     4406       +4     
==========================================
- Hits        96043    95741     -302     
- Misses      24147    24446     +299     
- Partials     1566     1570       +4     

☔ View full report in Codecov by Harness.
📢 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]

Correct fix, well aligned with the array-api discipline. The generic dpmodel reduction helpers were hard-coding NumPy ops (Tensor.astype, np.sum) and passing a NumPy dtype class into the active namespace, so they crashed/mis-typed on Torch -- exactly the pt_expt path via make_model. Resolving the energy precision through get_xp_precision and using xp.astype/xp.sum fixes it, and NumPy numerics are preserved (the int-vs-float64 mask divisor is value-identical). The three Torch regressions genuinely fail pre-fix (raw torch.Tensor has no .astype; the graph path rejects the NumPy dtype) and assert a float64 torch.Tensor result. Minor, non-blocking: only Torch is committed-tested here, but it's the representative namespace that exposes both bugs and the NumPy path stays covered by existing tests. LGTM.

Dismissing to re-review through the /code-review skill per process.

[wanghan-iapcm]

The fix itself is correct and complete (verified: get_xp_precision resolves the namespace-native dtype, NumPy numerics are preserved bit-identically, and the three new tests genuinely fail pre-fix). One blocking gap before merge: the mask is None reduction branches are untested in both files. Mutation testing confirms it -- reverting the fix on those exact lines while keeping the rest still passes all three tests, so a regression there would go undetected. Per the repo convention of covering every reachable branch (including both sides of a boolean), please add coverage for the no-mask intensive path (and ideally the graph extensive path). Requesting changes only for the test coverage; the code is good.

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @iProzd — 2 reviews on exact changed files (deepmd/dpmodel/model/transform_output.py).

No review request was made automatically.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

[wanghan-iapcm]

Both review comments are addressed with genuine regression tests, verified against HEAD. The dense intensive/no-mask branch (transform_output.py) and the graph n_node-fallback + extensive branches (edge_transform_output.py) are each now exercised by a Torch test that asserts the result stays a torch tensor in energy precision. Confirmed empirically that all three pre-fix idioms raise on a Torch tensor (Tensor.astype -> AttributeError, xp.astype with a NumPy dtype class -> TypeError, np.sum on a tensor -> TypeError), so every new test fails pre-fix and passes after -- closing the mutation-survival gap the original comments flagged. No residual numpy idioms remain in either transform; the fix is complete.

[njzjz]

Addressed the CodeRabbit nitpick in bb988aea: the Torch-dependent test module now imports pytest and uses module-level pytest.importorskip("torch"), so the entire module skips cleanly when Torch is unavailable while preserving all six test cases when it is installed.

Revalidated the current PR head: source/tests/common/dpmodel/test_transform_output.py passes (6 tests), and Ruff reports no issues for the file.

Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.