fix(jax): add hessian energy loss#5607
Conversation
njzjz
force-pushed
the
fix/dpmodel-energy-hessian-loss
branch
from
b56a253 to
ae453b0
Compare
|
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 (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthrough
ChangesHessian Loss Feature
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 OpenGrep (1.23.0)source/tests/common/dpmodel/test_loss_ener.py┌──────────────┐ �[32m✔�[39m �[1mOpengrep OSS�[0m [00.13][ERROR]: unable to find a config; path deepmd/dpmodel/loss/ener.py┌──────────────┐ �[32m✔�[39m �[1mOpengrep OSS�[0m [00.15][ERROR]: unable to find a config; path deepmd/jax/train/trainer.py┌──────────────┐ �[32m✔�[39m �[1mOpengrep OSS�[0m [00.14][ERROR]: unable to find a config; path 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: 2
🤖 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/dpmodel/loss/ener.py`:
- Around line 565-574: The Hessian entry in `DPTrainer.data_requirements` is
advertising the wrong tensor shape for `label_requirement`. Update the
`DataRequirementItem("hessian", ...)` definition in `ener.py` so it matches the
real on-disk Hessian layout used by the new loss path and tests, rather than the
current atomic `ndof=1` schema. If the dataset loader cannot yet consume the
full Hessian tensor, add the loader support first and keep `has_h` gated until
the contract is consistent.
In `@deepmd/jax/train/trainer.py`:
- Around line 120-121: The Hessian enablement in Trainer setup is unguarded, so
`self.model.enable_hessian()` can break for model/loss combinations like the JAX
`zbl` path that do not implement it. Update the `Trainer` logic to check that
the model actually exposes `enable_hessian` before calling it, using the
existing `self.loss.has_h` condition as the trigger and keeping the current
hessian tensor handling unchanged.
🪄 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: 23778160-e75a-4d2b-be1f-2babb0b6710f
📒 Files selected for processing (3)
deepmd/dpmodel/loss/ener.pydeepmd/jax/train/trainer.pysource/tests/common/dpmodel/test_loss_ener.py
| if self.has_h: | ||
| label_requirement.append( | ||
| DataRequirementItem( | ||
| "hessian", | ||
| ndof=1, | ||
| atomic=True, | ||
| must=False, | ||
| high_prec=False, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
label_requirement advertises the wrong Hessian shape.
DataRequirementItem("hessian", ndof=1, atomic=True) describes an nframes x natoms x 1 label, but the new loss path and tests consume a full nframes x (3 * natoms) x (3 * natoms) tensor. DPTrainer.data_requirements forwards this schema to the dataset loader, so real Hessian training will request/load hessian.npy with an incompatible contract even though these unit tests pass by building label_dict manually. Please make the requirement describe the real on-disk Hessian layout, or add dedicated loader support before enabling this in training.
🤖 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/dpmodel/loss/ener.py` around lines 565 - 574, The Hessian entry in
`DPTrainer.data_requirements` is advertising the wrong tensor shape for
`label_requirement`. Update the `DataRequirementItem("hessian", ...)` definition
in `ener.py` so it matches the real on-disk Hessian layout used by the new loss
path and tests, rather than the current atomic `ndof=1` schema. If the dataset
loader cannot yet consume the full Hessian tensor, add the loader support first
and keep `has_h` gated until the contract is consistent.
1 participant
Summary
EnergyLossinstead of introducing a separate loss classTests
source venv/bin/activate && pytest source/tests/common/dpmodel/test_loss_ener.py -qsource venv/bin/activate && ruff check .source venv/bin/activate && ruff format .Summary by CodeRabbit
New Features
Tests