fix(jax): materialize default fparam in cxx api#5849
Open
njzjz-bot wants to merge 1 commit into
Open
Conversation
Contributor
📝 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}
Contributor
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
njzjz-bot
force-pushed
the
fix/jax-default-fparam-cxx-api-5658
branch
from
July 17, 2026 01:46
d23eb87 to
ffd8a5c
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5849 +/- ##
==========================================
- Coverage 78.58% 78.30% -0.28%
==========================================
Files 1050 1050
Lines 120637 120753 +116
Branches 4356 4381 +25
==========================================
- Hits 94801 94557 -244
- Misses 24278 24616 +338
- Partials 1558 1580 +22 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
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 #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
Bug Fixes
Tests