fix(pt-expt): correct aparam reshape in graph mode

[MoseyQAQ]

Problem

Training with aparam in graph mode fails with a dimension mismatch.

Cause

The reshape uses ap.shape[-1] as the per-atom parameter dimension. For flattened inputs shaped (nframes, natoms * dim_aparam), this is natoms * dim_aparam, causing an invalid reshape.

Fix

Use self.get_dim_aparam() to reshape the input with the correct per-atom parameter dimension.

Summary by CodeRabbit

• Bug Fixes
  • Improved graph processing for atomic parameters with explicitly configured dimensions.
  • Fixed parameter reshaping and indexing for more reliable model behavior.
  • Ensured flattened per-atom parameter inputs produce results consistent with the standard layout across energy, force, and virial outputs in both training and evaluation modes.
  • Added validation that atomic parameter values correctly influence energy predictions.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Advanced

Run ID: ca69bee8-fa56-4fb1-b279-c3f8b724c38e

📥 Commits

Reviewing files that changed from the base of the PR and between 89ef1d6 and ec40275.

📒 Files selected for processing (1)

• source/tests/pt_expt/model/test_graph_export.py

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

_call_common_graph now reshapes ap with the declared atomic parameter dimension before compacted-node indexing. Tests verify equivalent flattened and canonical inputs across graph outputs in training and evaluation modes.

Changes

Atomic parameter reshape

Layer / File(s)	Summary
Use the declared atomic parameter width deepmd/pt_expt/model/make_model.py	The ap reshape uses self.get_dim_aparam() before indexing with compacted node indices.
Validate flattened atomic parameters source/tests/pt_expt/model/test_graph_export.py	Tests compare flattened and canonical aparam inputs across energy, force, virial, and atomic outputs in training and evaluation modes.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Change: Bug fix

Suggested reviewers: wanghan-iapcm

Merge Risk: ⚪ Minimal · up to ec402

The atomic-parameter reshape correction is covered for the affected flattened graph inputs, with no actionable merge-blocking risk identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: correcting the aparam reshape in graph mode.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

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

[njzjz-bot]

The implementation direction looks correct: reshaping by the model-declared per-atom parameter width fixes the reported (nframes, natoms * dim_aparam) graph-mode layout while preserving the regular flat-node ABI. However, this exact regression is still not covered by a test. The existing test_graph_export_aparam_flat_node_axis supplies (N, dim_aparam), where the old reshape already works, so it would not have caught this failure. Please add focused regression coverage for the failing graph-mode/training layout before merging. Relevant exact-head CI is also still incomplete (Test Python, Test C++, and CodeQL were still running at review time).

— Agent: ChatGPT · Model: GPT-5.6 Sol · GitHub account: njzjz-bot · Reviewed head: 89ef1d6 · Trigger: scheduled all-PR monitoring

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.11%. Comparing base (484816a) to head (ec40275).
⚠️ Report is 4 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #6028      +/-   ##
==========================================
- Coverage   77.29%   77.11%   -0.18%     
==========================================
  Files        1154     1151       -3     
  Lines      139211   139000     -211     
  Branches     5056     5056              
==========================================
- Hits       107598   107195     -403     
- Misses      29727    29921     +194     
+ Partials     1886     1884       -2     

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

[njzjz-bot]

The previously requested regression coverage is now addressed: the new test exercises the actual failing (nframes, natoms * dim_aparam) layout with multiple frames/atoms/components in both training and evaluation modes, compares against the canonical per-atom layout, and checks the relevant energy/force/virial outputs. I found no remaining high-confidence correctness issue in the two-file change.

I am not approving yet because the exact-head GitHub Actions workflows did not actually execute: Build C library, CodeQL, Test C++, Test Python, Test CUDA, PyPI/package, and Build C++ are all action_required, and the inspected Test Python run has zero jobs. This looks like the external-fork workflow approval gate rather than a test failure. Once those exact-head checks are allowed to run and pass, this should be ready for approval.

— Agent: ChatGPT · Model: GPT-5.6 Sol · GitHub account: njzjz-bot · Reviewed head: ec40275 · Trigger: scheduled all-PR monitoring

[MoseyQAQ]

@OutisLi Hi, Can you help check it?

[njzjz-bot]

The previously requested regression coverage is present on this head and directly exercises the failing (nframes, natoms * dim_aparam) layout in both training and evaluation modes, comparing it against the canonical per-atom layout across energy, force, virial, and atomic outputs. The implementation itself is the minimal correct fix: reshape by the model-declared per-atom width before compact-node indexing.

The exact-head CI has now completed successfully: Test Python, Test C++, Test CUDA, Build C++, Build C library, CodeQL, and PyPI/package all passed. I found no remaining high-confidence correctness, API, or test-coverage blocker.

— Agent: ChatGPT · Model: GPT-5.6 Sol · GitHub account: njzjz-bot · Reviewed head: ec40275 · Trigger: scheduled all-PR monitoring