fix(infer): normalize parameter shorthand before batching#5857
Open
njzjz-bot wants to merge 2 commits into
Open
fix(infer): normalize parameter shorthand before batching#5857njzjz-bot wants to merge 2 commits into
njzjz-bot wants to merge 2 commits into
Conversation
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (8)
🚧 Files skipped from review as they are similar to previous changes (6)
📝 WalkthroughWalkthroughDeepEval gains shared validation and expansion for ChangesParameter shorthand normalization
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related issues
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/pt/infer/deep_eval.py`:
- Around line 548-555: Update eval_embedding to call _standardize_fparam_aparam
on fparam and aparam before dispatching to self._eval_func, using the embedding
evaluation’s frame count, natoms, and model parameter dimensions. Pass the
standardized values to the backend so shorthand inputs are batched by frames
correctly.
🪄 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
Run ID: e2d24567-276c-4ec1-a33e-25ee33258076
📒 Files selected for processing (7)
deepmd/infer/deep_eval.pydeepmd/jax/infer/deep_eval.pydeepmd/pd/infer/deep_eval.pydeepmd/pt/infer/deep_eval.pydeepmd/tf2/infer/deep_eval.pysource/tests/common/test_deep_eval_parameter_shorthand.pysource/tests/consistent/io/test_io.py
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
njzjz-bot
force-pushed
the
fix/deepeval-param-shorthand-backends-5666
branch
from
July 17, 2026 01:57
2149d24 to
9227907
Compare
for more information, see https://pre-commit.ci
Member
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #5666.
Summary
fparamandaparamsize checks and tiling rulesAutoBatchSizein the PyTorch, JAX, TF2, and Paddle backend adapterseval_embeddingbefore batching, which also fixes itseval_descriptorandeval_fitting_last_layerwrappersWhy existing tests missed this
Existing multi-frame inference tests call the high-level
DeepPot.evalwrapper, whose_standard_inputalready expands shorthand before backend dispatch. Backend-direct consistency tests used canonical full shapes, while many backend-specific fixtures use only one frame, where shorthand sizes can coincide with full sizes. Spin inference fixtures also use zero-dimensionalfparam/aparam. TF2/JAX coverage runs only in the TF2-only job, and Paddle was absent from the cross-backend IO matrix.The regressions deliberately call
.deep_eval.evaland the PyTorch backend's.eval_embeddingto bypass public normalization, supplies(dim_fparam,),(natoms, dim_aparam), and(dim_aparam,), and forcesauto_batch_size=natomsso each frame is evaluated separately. This catches both direct reshape failures and accidental slicing of the atom axis.Validation
pytest source/tests/common/test_deep_eval_parameter_shorthand.py -q(7 passed)TestDeepPotFparamAparam.test_deep_evalwith backend-direct forced batching (passed)TestEmbeddingDeepEvalAPI(6 passed), including shared fparam, per-atom aparam, and scalar aparam through forced one-frame batchingruff format .(1664 files unchanged on final pass)ruff check .(passed)git diff --check(passed)The pre-existing optional PyTorch C++ extension still has an ABI mismatch with the current Torch package, so local tests disabled loading that unrelated
.so; the pure PyTorch backend-direct embedding suite ran successfully.Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
New Features
Bug Fixes
Tests