fix(dpmodel): support parameter shorthand in DeepEval - #5853
Conversation
Standardize frame and atomic parameter shorthand in the common wrapper and backend entry points before automatic batching. Preserve full per-frame and per-atom arrays while broadcasting documented shared forms consistently across TensorFlow 2, PyTorch, JAX, and Paddle. Normalize PyTorch embedding extraction before split execution so eval_embedding, eval_descriptor, and eval_fitting_last_layer accept shared fparam, per-atom aparam, and scalar aparam forms. Cover two frames forced through one-frame backend batches. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
for more information, see https://pre-commit.ci
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5853 +/- ##
==========================================
+ Coverage 79.03% 79.24% +0.21%
==========================================
Files 1055 1072 +17
Lines 122233 125055 +2822
Branches 4401 4541 +140
==========================================
+ Hits 96607 99103 +2496
- Misses 24061 24325 +264
- Partials 1565 1627 +62 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
There was a problem hiding this comment.
Fix is correct. AutoBatchSize.execute_all slices only ndim > 1 args along axis 0, so a (natoms, dim_aparam) shorthand really was being sliced on its atom axis - canonicalizing before the batcher is the right place. I checked the size-dispatch for ambiguity: collisions occur only at n == 1 or natoms == 1, and in each the two branches produce the identical array, so the check order is safe. Predicate order matches _standard_input, so the public path stays idempotent with no double-tiling. All 8 new tests fail pre-fix, verified against an unpatched tree.
Two comments inline.
Merge ordering with #5852. It inserts at this same anchor and rewrites the out = self._eval_func(...) call this diff carries as context; git merge-tree confirms a real conflict, not just adjacency. Both PRs exist for the same reason - AutoBatchSize slices only frame-axis arrays - and spin is likewise unnormalized in _standard_input. Canonicalizing in the shared layer would cover spin, fparam and aparam together and collapse both PRs into one change.
Minor:
-
_standard_inputemits aparam as 2-D(nframes, natoms * dim_aparam); this returns 3-D._eval_modelreshapes either way, but it changes the internal boundary #4285 established without documenting it. -
AutoBatchSize.__init__takescurrent_batch_sizefromDP_INFER_BATCH_SIZEwhen set, overriding the constructor argument. Under that env var theauto_batch_size=6case stops splitting and silently becomes a duplicate ofFalse, so the only batching coverage disappears with no failure.monkeypatch.delenv("DP_INFER_BATCH_SIZE", raising=False)plus an assertion on the_eval_modelcall count would pin it. -
Validation is size-only, so a transposed
(dim_fparam, nframes)array passes and is reinterpreted row-major. Matches_standard_input, so not a regression, but it sits against the stated goal of explicit shape errors.
| natoms, numb_test = self._get_natoms_and_nframes( | ||
| coords, atom_types, len(atom_types.shape) > 1 | ||
| ) | ||
| # Canonicalize documented shorthand before AutoBatchSize sees it. |
There was a problem hiding this comment.
This ladder already exists in DeepEval._standard_input (deepmd/infer/deep_eval.py) - same three-way dispatch, same tiling, same wording. That placement was deliberate: #3213 put canonicalization in _standard_input, #3356 added this ladder there for every backend, and #4285 gave this adapter only the bare reshape because it consumes already-rectangular arrays.
The gap isn't dpmodel-specific. All three siblings reshape without tiling:
deepmd/pt/infer/deep_eval.py:695deepmd/jax/infer/deep_eval.py:411deepmd/pd/infer/deep_eval.py:544
while DeepEvalBackend.eval's docstring - the contract all of them implement - documents all three shorthands. So after this PR deep_eval.eval(fparam=<1-D>) succeeds on .dp and raises on .pth.
I'd rather see this fixed once in the shared layer than per backend: a DeepEvalBackend._canonicalize_params(fparam, aparam, nframes, natoms) called from each backend's eval covers all five at once. As written this is the third copy, and the copies already disagree - _standard_input says "frame param" where this says "atomic param", so the same bad input reports differently depending on entry point.
There was a problem hiding this comment.
Fixed by stacking this PR on #5857 and force-updating the branch to f340567. The dpmodel adapter now uses the single shared _standardize_fparam_aparam helper from #5857; the duplicated inline ladder has been removed. This PR now contains only the focused dpmodel regression on top of the shared backend fix.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
| ).reshape(len(ATOM_TYPES), 3) | ||
| BOX = np.diag([13.0, 13.0, 13.0]) | ||
| NFRAMES = 2 | ||
| COORDS = np.tile(COORD, (NFRAMES, 1, 1)) |
There was a problem hiding this comment.
COORDS = np.tile(COORD, (NFRAMES, 1, 1)) makes both frames byte-identical, and every test's parameters are frame-invariant, so nothing here distinguishes correct frame-major handling from "use frame 0's parameters everywhere".
The tests do fail pre-fix, but as a crash (cannot reshape array of size 2 into shape (1,6,2)), not as a wrong number - a regression that kept the shapes right while losing per-frame variation would pass all 8.
One case closes it: fparam (NFRAMES, dim) with distinct rows and aparam (NFRAMES, natoms, dim) with distinct frames, split vs unsplit, asserting both that the two agree and that the per-frame energies differ. The second assertion is what makes the first meaningful.
There was a problem hiding this comment.
Fixed in f340567. The fixtures now use distinct coordinates and boxes, and a dedicated case supplies distinct fparam and aparam rows per frame. It evaluates the split path first on a fresh batcher, compares against an unsplit evaluator, and asserts the two frame energies differ.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
…shorthand-backends-5666
deserialize_to_file only writes .json, serialize_from_file raises NotImplementedError, and the .json reader rejects fparam/aparam, so the paddle entry could not complete the round trip. Its normalization stays covered by test_deep_eval_parameter_shorthand.py.
Normalize parameter shorthand in the remaining dpmodel, pt_expt, and TensorFlow adapters before automatic batching. Strengthen the regressions with distinct frames and parameters, fresh batchers, and direct Paddle and pt_expt adapter coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
1eb3feb to
f340567
Compare
|
Warning Review limit reached
Next review available in: 40 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (6)
📝 WalkthroughWalkthroughDeepEval now normalizes shared and full ChangesParameter shorthand support
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant DeepEval
participant StandardizeFparamAparam
participant ModelEvaluator
DeepEval->>StandardizeFparamAparam: pass frame count, atom count, and parameter dimensions
StandardizeFparamAparam-->>DeepEval: return normalized fparam and aparam
DeepEval->>ModelEvaluator: evaluate with normalized parameters
Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/infer/deep_eval.py (1)
992-1032: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the
_standard_inputreturn type hint to match its actual 7-value return.The declared return type is a 5-element tuple, but the function returns
coords, cells, atom_types, fparam, aparam, nframes, natoms— 7 values. Callers such aseval_descriptoralready unpack all 7 values, so the annotation is stale. Addint, intfornframesandnatomsto the return type.🔧 Proposed fix for the return type hint
def _standard_input( self, coords: np.ndarray | list, cells: np.ndarray | list | None, atom_types: np.ndarray | list, fparam: np.ndarray | list | None, aparam: np.ndarray | list | None, mixed_type: bool, ) -> tuple[ np.ndarray, np.ndarray | None, np.ndarray, np.ndarray | None, np.ndarray | None, + int, + int, ]:🤖 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 `@deepmd/infer/deep_eval.py` around lines 992 - 1032, Update the return type annotation of _standard_input to include the two trailing int values for nframes and natoms, matching its seven-value return tuple and existing caller unpacking.
🧹 Nitpick comments (1)
source/tests/pd/test_deep_eval_parameter_shorthand.py (1)
21-55: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider extracting the shared dispatch-normalization test.
This test duplicates
test_eval_standardizes_parameter_shorthand_before_dispatchfromsource/tests/pt_expt/infer/test_parameter_shorthand.pyalmost line for line, differing only in the importedDeepEvalmodule. Extract a shared, backend-parametrized helper (for example, a small factory that accepts the module path or the class) to avoid copy-paste drift as more backends adopt this dispatch check.🤖 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/pd/test_deep_eval_parameter_shorthand.py` around lines 21 - 55, Extract the duplicated setup and assertions from test_eval_standardizes_parameter_shorthand_before_dispatch into a shared backend-parameterized helper or factory that accepts the relevant DeepEval class/module. Update both this test and the corresponding test_parameter_shorthand.py test to invoke the shared helper while preserving each backend’s existing evaluator behavior and assertions.
🤖 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.
Inline comments:
In `@deepmd/tf/infer/deep_eval.py`:
- Around line 765-772: Move the _standardize_fparam_aparam call to the beginning
of eval_descriptor, before any auto-batching or _eval_func dispatch. Normalize
shorthand and shared fparam/aparam forms once using the existing frame,
atom-count, and dimension arguments, ensuring shared (natoms, dim_aparam)
parameters are not split along the frame axis.
---
Outside diff comments:
In `@deepmd/infer/deep_eval.py`:
- Around line 992-1032: Update the return type annotation of _standard_input to
include the two trailing int values for nframes and natoms, matching its
seven-value return tuple and existing caller unpacking.
---
Nitpick comments:
In `@source/tests/pd/test_deep_eval_parameter_shorthand.py`:
- Around line 21-55: Extract the duplicated setup and assertions from
test_eval_standardizes_parameter_shorthand_before_dispatch into a shared
backend-parameterized helper or factory that accepts the relevant DeepEval
class/module. Update both this test and the corresponding
test_parameter_shorthand.py test to invoke the shared helper while preserving
each backend’s existing evaluator behavior and assertions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 618e92dd-e067-4b32-a679-b13869cc5abb
📒 Files selected for processing (14)
deepmd/dpmodel/infer/deep_eval.pydeepmd/infer/deep_eval.pydeepmd/jax/infer/deep_eval.pydeepmd/pd/infer/deep_eval.pydeepmd/pt/infer/deep_eval.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/tf/infer/deep_eval.pydeepmd/tf2/infer/deep_eval.pysource/tests/common/test_deep_eval_parameter_shorthand.pysource/tests/consistent/io/test_io.pysource/tests/infer/test_dpmodel_deep_eval_params.pysource/tests/pd/test_deep_eval_parameter_shorthand.pysource/tests/pt/model/test_embedding.pysource/tests/pt_expt/infer/test_parameter_shorthand.py
Coding-Agent: Codex\nCodex-Version: codex-cli 0.144.6\nModel: gpt-5.6-sol\nReasoning-Effort: xhigh
f340567 to
e7505f2
Compare
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both of my earlier points are resolved, and the first one in the better way -- rather than adding a third copy of the ladder, the branch is now stacked on #5857 so the dpmodel adapter calls the single shared _standardize_fparam_aparam. That also removes the message divergence I complained about: master's dpmodel copy reported "wrong size of frame param" for an aparam error, and the shared helper gets that right.
On the second point, the fixtures now vary: COORDS = np.stack((COORD, COORD + 0.125)), BOXES[1, 0, 0] += 0.25, and test_distinct_full_parameters_preserve_frame_order_when_split supplies distinct fparam and aparam rows per frame and asserts the two frame energies differ. That last assertion is the one that makes the comparison meaningful, and it is here.
I ran the file against pre-fix source rather than reading it: 8 of its 9 cases fail, with real diagnostics (cannot reshape array of size 2 into shape (2,6,2) at deepmd/dpmodel/infer/deep_eval.py:366, and both error-message cases failing because master raises a bare ValueError instead of the worded RuntimeError). The 9th is test_distinct_full_parameters_preserve_frame_order_when_split, which passes pre-fix -- correctly so, since it feeds full frame-major arrays that never needed tiling. Its value is as a guard against a future change that keeps the shapes right while losing per-frame variation, which is exactly what I asked for, so I want to be clear I am not counting it against the PR.
Approving. Two notes, neither blocking.
Merge order. This branch carries #5857's commits, so landing it first would merge that PR's contents too. I have an open point on #5857 about eval_descriptor / eval_fitting_last_layer in the pt_expt adapter still bypassing the shared helper; that is worth settling before either lands, since merging this one silently decides it.
Worth promoting upward. Three things this file does that #5857's shared _assert_backend_parameter_shorthand does not, called out inline -- they would help every backend rather than dpmodel alone.
| ) | ||
|
|
||
| _assert_outputs_equal(split, unsplit) | ||
| assert not np.allclose(split["energy_redu"][0], split["energy_redu"][1]) |
There was a problem hiding this comment.
This assertion is the most valuable line in the file, and it exists only here.
_assert_backend_parameter_shorthand in #5857 compares shorthand against a full frame-major reference for tensorflow, pytorch and dpmodel, but never asserts that the two frames actually differ -- so a degenerate result where both frames came out identical would satisfy it. That is the same gap I originally raised against this file, now surviving in the shared helper that covers more backends.
Two smaller ones in the same direction:
- This file parametrizes
auto_batch_sizeover[False, len(ATOM_TYPES)], exercising both the split and unsplit paths. The shared helper only ever runs the split path. test_invalid_parameter_size_has_clear_errorasserts theRuntimeErrorpropagates out of a realDeepEval.evalcall. fix(infer): normalize parameter shorthand before batching #5857 tests the error branches only by calling_standardize_fparam_aparamdirectly insource/tests/common/test_deep_eval_parameter_shorthand.py, so nothing there proves the exception survives the adapter and the batcher.
None of this needs changing in this PR -- the coverage is correct as written. But when #5857 merges and this rebases down to just this file, folding these three into the shared helper would give every backend what dpmodel gets here, instead of leaving dpmodel with the stronger suite by accident.
There was a problem hiding this comment.
Agreed — this assertion stays here, and the same gap is now closed in the shared helper. _assert_backend_parameter_shorthand in #5857's source/tests/consistent/io/test_io.py now asserts that the two input frames produce different per-frame energy values on every shorthand case (commit 1709902 on #5857), so a degenerate expansion that reuses frame 0's parameters everywhere can no longer satisfy the comparison by also making the reference degenerate. This PR was rebased onto the updated #5857 (692789d); its own test still passes (9 tests).
Coding agent: opencode
opencode version: 1.18.9
Model: ustc/deepseek-v4-flash
Reasoning effort: max
…r routes Wire _standardize_fparam_aparam into eval_descriptor and eval_fitting_last_layer so shared per-atom shorthand cannot be mistaken for a frame axis before _prepare_nlist_inputs reshapes it. Extend the pt_expt adapter regressions to cover both routes. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
The backend-direct shorthand comparison never checked that the two input frames actually differ, so a degenerate expansion that reuses one frame's parameters for every frame would also make the frame-major reference degenerate and slip past the allclose comparison. Assert the per-frame energy values differ on every shorthand case. Coding-Agent: opencode opencode-Version: 1.18.9 Model: ustc/deepseek-v4-flash Reasoning-Effort: max
Stack the dpmodel-specific regression on the shared backend normalization from deepmodeling#5857. Use distinct coordinates and frame-major parameters, run shorthand inputs on fresh batchers, and verify split evaluation preserves frame order and distinct energies. Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
e7505f2 to
692789d
Compare
3 participants
Closes #5662.
Depends on #5857, which provides the shared parameter-normalization helper and wires it into every backend. This PR is intentionally stacked on that branch and adds only dpmodel-specific regression coverage; after #5857 merges, the shared implementation commits will disappear from this PR's diff.
Summary
Validation
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh