feat(jax): freeze models with hessian output#5608
Open
njzjz wants to merge 6 commits into
Open
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 (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds a ChangesJAX freeze Hessian support
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant CLI
participant freeze
participant serialization
participant model
participant HLO
CLI->>freeze: freeze --hessian
freeze->>serialization: deserialize_to_file(hessian=True)
serialization->>model: enable_hessian()
serialization->>serialization: set model_def_script hessian_mode
serialization-->>HLO: export Hessian-enabled model
Possibly related PRs
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.
🧹 Nitpick comments (1)
source/tests/jax/test_training.py (1)
199-214: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert
hessianin the dispatcher call.This still passes if
main()dropsargs.hessianbefore invokingfreeze(). Please assert the patched call carries the new flag as well.Suggested test tightening
main(args) freeze_entrypoint.assert_called_once() + self.assertIn("hessian", freeze_entrypoint.call_args.kwargs) + self.assertFalse(freeze_entrypoint.call_args.kwargs["hessian"])🤖 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/jax/test_training.py` around lines 199 - 214, The test for the JAX CLI freeze dispatch is too loose because it only checks that the patched freeze entrypoint was called. Tighten test_main_dispatches_freeze in test_training.py by asserting the call includes the hessian flag from the argparse.Namespace, so main() must forward args.hessian into deepmd.jax.entrypoints.main.freeze rather than dropping it. Use the existing freeze_entrypoint mock and verify its call arguments reflect the new flag.
🤖 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.
Nitpick comments:
In `@source/tests/jax/test_training.py`:
- Around line 199-214: The test for the JAX CLI freeze dispatch is too loose
because it only checks that the patched freeze entrypoint was called. Tighten
test_main_dispatches_freeze in test_training.py by asserting the call includes
the hessian flag from the argparse.Namespace, so main() must forward
args.hessian into deepmd.jax.entrypoints.main.freeze rather than dropping it.
Use the existing freeze_entrypoint mock and verify its call arguments reflect
the new flag.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 44d1886d-fd38-4bbe-a056-c56edc3c92fc
📒 Files selected for processing (7)
deepmd/jax/entrypoints/freeze.pydeepmd/jax/infer/deep_eval.pydeepmd/jax/jax2tf/serialization.pydeepmd/jax/model/hlo.pydeepmd/jax/utils/serialization.pydeepmd/main.pysource/tests/jax/test_training.py
njzjz
force-pushed
the
feat/jax-hessian-freeze
branch
from
16058df to
6eabd63
Compare
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5608 +/- ##
==========================================
- Coverage 79.65% 79.52% -0.14%
==========================================
Files 1015 1015
Lines 115740 115805 +65
Branches 4275 4272 -3
==========================================
- Hits 92198 92096 -102
- Misses 22001 22167 +166
- Partials 1541 1542 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Member
Author
|
I'd like to wait for #5613. |
njzjz
force-pushed
the
feat/jax-hessian-freeze
branch
from
a476689 to
30735c1
Compare
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Avoid requesting Hessian output from standard frozen JAX models and add regression coverage. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
wanghan-iapcm
approved these changes
Jul 12, 2026
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
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.
Summary
dp freeze --hessiansupport for the JAX backend.jax,.hlo, and.savedmodelserializationTests
source venv/bin/activate && pytest source/tests/jax/test_training.py::TestJAXTraining::test_freeze_entrypoint_uses_checkpoint_pointer source/tests/jax/test_training.py::TestJAXTraining::test_main_dispatches_freeze source/tests/jax/test_training.py::TestJAXTraining::test_hlo_hessian_mode_updates_output_def source/tests/jax/test_training.py::TestJAXTraining::test_deep_eval_requests_hessian_for_hessian_model -qsource venv/bin/activate && ruff check .source venv/bin/activate && ruff format .Summary by CodeRabbit
New Features
--hessianoption to the model freezing command to include Hessian information in exported outputs.Bug Fixes
Tests