fix(jax): materialize default fparam in cxx api - #5849
Conversation
|
Warning Review limit reached
Next review available in: 31 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 (4)
📝 WalkthroughWalkthroughDeepPotJAX now loads embedded default frame parameters, validates tensor shapes and sizes, normalizes explicit or default ChangesJAX frame-parameter handling
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant DeepPotJAX
participant JAXSavedModel
participant make_fparam_input
participant create_tensor
DeepPotJAX->>JAXSavedModel: load model metadata and default_fparam
JAXSavedModel-->>DeepPotJAX: return default values
DeepPotJAX->>make_fparam_input: normalize supplied or default fparam
make_fparam_input->>create_tensor: create validated model input tensor
create_tensor-->>JAXSavedModel: provide fparam tensor for inference
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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
🧹 Nitpick comments (1)
source/api_cc/tests/test_deeppot_universal.cc (1)
1994-2050: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover the neighbor-list path in the JAX-specific regressions.
These tests exercise multiframe broadcasting and invalid-size rejection only through direct
dp.compute. Add equivalent neighbor-list cases so regressions in that separate overload cannot pass unnoticed.🤖 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/api_cc/tests/test_deeppot_universal.cc` around lines 1994 - 2050, Extend the JAX-specific regression coverage in DefaultFParamDeepPotTest by adding neighbor-list overload cases equivalent to JAXBroadcastsFParamAcrossFrames and JAXRejectsInvalidFParamSize. Use the existing multiframe/default-versus-override expectations and invalid fparam-size assertion, but invoke the neighbor-list compute path so both broadcasting and rejection are validated there.
🤖 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 `@source/api_cc/src/DeepPotJAX.cc`:
- Around line 443-469: Update the tensor creation code around TF_AllocateTensor
to cast shape.size() explicitly to the expected integer type, then check whether
TF_AllocateTensor returned nullptr before accessing TF_TensorData. Throw a
deepmd::deepmd_exception with an allocation-failure message when allocation
fails, while preserving the existing copy behavior for successful allocations.
---
Nitpick comments:
In `@source/api_cc/tests/test_deeppot_universal.cc`:
- Around line 1994-2050: Extend the JAX-specific regression coverage in
DefaultFParamDeepPotTest by adding neighbor-list overload cases equivalent to
JAXBroadcastsFParamAcrossFrames and JAXRejectsInvalidFParamSize. Use the
existing multiframe/default-versus-override expectations and invalid fparam-size
assertion, but invoke the neighbor-list compute path so both broadcasting and
rejection are validated there.
🪄 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: 489a8196-07d7-4e0e-a644-db1f89845485
📒 Files selected for processing (4)
source/api_cc/include/DeepPotJAX.hsource/api_cc/src/DeepPotJAX.ccsource/api_cc/tests/test_deeppot_universal.ccsource/tests/infer/gen_fparam_aparam.py
Load and validate default frame parameters from JAX SavedModels, materialize stored or caller-provided values for direct and neighbor-list inference, and broadcast one-frame values across multiple frames. Validate TensorFlow C API tensor shapes and sizes before allocation, cast rank explicitly, and report allocation failure before accessing tensor storage. Cover default and override behavior, float and double paths, multiframe broadcasting, neighbor lists, and invalid sizes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
d23eb87 to
ffd8a5c
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5849 +/- ##
==========================================
+ Coverage 79.03% 79.21% +0.17%
==========================================
Files 1055 1072 +17
Lines 122233 125169 +2936
Branches 4401 4566 +165
==========================================
+ Hits 96607 99148 +2541
- Misses 24061 24384 +323
- Partials 1565 1637 +72 ☔ 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 |
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The fix looks right, and it closes a genuine C++/Python divergence — #5311 wired the has_default_fparam flag through all the C++ backends and #5343 implemented substitution for .pt2 only, so JAX has been carrying the flag without the behaviour since. The precedence (caller fparam > stored default > error) matches deepmd/jax/infer/deep_eval.py, the single-frame broadcast arithmetic is correct, and I specifically checked that older SavedModels without get_default_fparam still load: the tf_function_not_found catch warns rather than throwing, matching the graceful-degradation pattern from 3e779f1. The explicit-override case (0.5 against a stored 0.25852028) is a good choice — it would catch the opposite bug too.
Two notes inline, plus one thing worth a separate issue rather than this PR: aparam's single-frame broadcast is documented in DeepPotJAX.h ("then all frames are assumed to be provided with the same aparam") but never implemented, so a multi-frame call with single-frame aparam reads past the input in deepmd::select_map — an out-of-bounds read, not just wrong values. The new JAXBroadcastsFParamAcrossFrames test pre-tiles aparam by hand (repeat_values(deepmd_test::aparam_value(), nframes)), which is exactly what sidesteps it. Pre-existing and out of scope here, but now that fparam and charge/spin both broadcast, aparam is the only per-frame input left without it.
Coding-Agent: Codex Codex-Version: codex-cli 0.144.6 Model: gpt-5.6-sol Reasoning-Effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both points are addressed.
The coverage gap is closed properly rather than minimally: jax_savedmodel is now in fparam_aparam_cases() alongside the existing default_fparam_cases() entry, and gen_fparam_aparam.py exports both fparam_aparam_default.savedmodel and fparam_aparam.savedmodel. That gives JAX a fixture on each side of has_default_fparam, so the "required but not provided, and no default_fparam is stored" throw is finally reachable for the backend this PR changes. The generation is wired into source/install/test_cc_local.sh, so the fixtures exist where the tests look for them.
On the second point, make_fparam_input now warns to stderr and returns an empty vector when dim_fparam == 0 and a non-empty fparam arrives, which is what make_charge_spin_input already did and what Python does with a UserWarning. The two helpers agree now, which was all I wanted -- I did not mind which way it went, only that the same user mistake not produce a warning in one place and an abort in the other.
One thing worth saying plainly about validation, because the check marks are misleading here. Of the four Test C++ jobs, one failed on the Paddle inference library download and two were cancelled by fail-fast; the only one that completed configures with ENABLE_TENSORFLOW=FALSE -D ENABLE_PYTORCH=FALSE -D ENABLE_PADDLE=TRUE. source/api_cc/CMakeLists.txt builds src/DeepPotJAX.cc only inside if(ENABLE_JAX), and ENABLE_JAX defaults off and is switched on by ENABLE_TENSORFLOW. So no CI job at this head compiled the file this PR changes, let alone ran the new JAX cases. Nothing about that is your doing -- it is the Paddle download outage cascading -- but it does mean the new coverage is currently unexecuted, and the value of these two fixtures is precisely that they run.
Approving on review of the code and the fixture wiring. Worth confirming a TensorFlow-enabled Test C++ job goes green before this merges, so the JAX rows actually prove something.
5902263
3 participants
Closes #5658.
Summary
get_default_fparamwhen initializing a JAX SavedModel in the C++ APIWhy existing tests missed this
The shared C++ default-fparam suite only included TorchScript and PT2 artifacts. Existing JAX C++ fixtures either had
dim_fparam == 0or supplied frame parameters explicitly, while Python JAX inference materializes the saved default before invoking the model. Serialization coverage therefore proved that the flag/getter were exported, but never exercised C++ consumption of an omitted JAXfparam.The explicit regression now uses
fparam=[0.5], distinct from the stored default[0.25852028], so it also detects an implementation that ignores caller overrides.Validation
ruff format .ruff check .clang-format --dry-run --Werroron all changed C++ filesrunUnitTests_ccanddeepmd_backend_jax; both targets rebuilt successfully after the allocation-safety follow-upDefaultFParamDeepPotTestJAX SavedModel cases passed, covering direct and neighbor-list inference, float/double, stored defaults, distinct explicit overrides, two-frame broadcasting, metadata, and invalid sizesCoding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit