[pull] master from deepmodeling:master - #265
Merged
Merged
Conversation
## Summary - add `dp freeze --hessian` support for the JAX backend - propagate the Hessian flag through JAX `.jax`, `.hlo`, and `.savedmodel` serialization - mark HLO energy output definitions as Hessian-enabled and request Hessian outputs during JAX inference ## Tests - `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 -q` - `source venv/bin/activate && ruff check .` - `source venv/bin/activate && ruff format .` <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added a `--hessian` option to the model freezing command to include Hessian information in exported outputs. * Evaluation now detects Hessian-capable models and can request Hessian-related derivative outputs when enabled. * **Bug Fixes** * Preserved Hessian mode during export so the frozen model metadata and output definitions correctly reflect Hessian settings. * Updated energy output handling to use Hessian-aware energy definitions when Hessian mode is active. * **Tests** * Added/updated regression and CLI tests to verify Hessian flag propagation through freezing and correct Hessian-aware output selection in evaluation. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn> Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>
## Summary Enable eight Ruff rules that have safe, automatic fixes in the repository's pinned Ruff v0.15.18: - `PIE808`, `PT022`, `RET502`, `SIM910`, `TD006` - `PLR1711`, `PLR1733`, `PLR2044` ## Intentional follow-up This PR deliberately changes only `pyproject.toml`; it does **not** include the 128 existing code fixes. The configured pre-commit.ci autofix should apply those fixes to this PR, which provides an end-to-end confirmation that the automated repair is clean. `COM812` is intentionally excluded because `ruff-format` is already enabled and Ruff recommends against combining the formatter with this trailing-comma rule. Rules whose fixes are classified as unsafe in Ruff v0.15.18 are also excluded. ## Validation - `git diff --check` - `ruff v0.15.18 check --output-format json .` reports exactly 128 intentional diagnostics from the eight newly enabled rules; no `--fix` was run. Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.6-terra) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved attribute assignment handling to fully control weight/bias-related updates during tracing. - Ensured neighborhood-distance resolution is initialized from available runtime/tensor sources and saved consistently. - Improved test isolation by clearing leaked device contexts before tests. - **Refactor** - Simplified option lookups (optional fields/seed/precision mappings) and equivalent loop-bound expressions across models/kernels. - **Chores** - Expanded linting rules and removed stray placeholder comments/formatting artifacts across docs and code. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
## What - Add an `edge_norm` descriptor option to control normalization on feature branches that vanish at the cutoff, with matching PyTorch and DPModel implementations. - When disabled, remove hidden radial-MLP RMSNorm, bypass environment-seed FiLM and cross-focus competition norms, and use a unit variance floor for post-SO(2) residual scaling. - Serialize the new option and its nested module settings, expose it through input validation, and enable the cutoff-smooth mode in the DPA4 water example. - Add cutoff-smoothness, serialization, configuration-propagation, and PyTorch/DPModel parity coverage. ## Why DPA4 radial features carry a C3 cutoff envelope and should decay to zero at `rcut`. Per-edge RMSNorm divides out that envelope until the feature variance reaches the norm epsilon floor, which can introduce a narrow, non-physical curvature and force feature just inside the cutoff. The same small-signal amplification applies to other normalization sites driven by cutoff-vanishing inputs. The new mode preserves the cutoff envelope: FiLM modulation returns smoothly to identity, cross-focus weights approach uniform, and post-SO(2) residual messages remain linear near zero. ## Impact - `edge_norm=True` remains the default, so existing configurations retain their current behavior. - Setting `edge_norm=False` opts into the recommended cutoff-smooth behavior. - PyTorch and DPModel serialization and numerical behavior remain aligned in both modes. - Normalization on persistent node states and FFN branches is unchanged. ## Checks - Added a near-cutoff radial sweep regression that verifies the curvature spike is suppressed when radial RMSNorm is disabled. - Added structure and serialization round-trip coverage for both normalization modes. - Added block- and descriptor-level PyTorch/DPModel parity coverage for standard and unit-floor post-SO(2) scaling. - Tests were run separately by the author; they were not rerun in this checkout per request. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added `edge_norm` for DPA4 and SeZM descriptors to control edge-dependent normalization behavior. * Added `radial_norm` for radial embeddings and `focus_norm` for SO(2) focus competition normalization. * Added `so2_post_norm_eps` to tune SO(2 post-normalization stability. * Updated model save/restore so these normalization settings are preserved. * **Documentation / Examples** * Updated the water DPA4 example input to set `edge_norm` explicitly. * **Tests** * Added tests covering the new switches (including serialization) and parity/consistency across implementations. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
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 subscribe to this conversation on GitHub.
Already have an account?
Sign in.
3 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.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )