perf(dpa4): opt so3grid#5517
Conversation
There was a problem hiding this comment.
Pull request overview
This PR optimizes the SO(3) grid-net quadratic operations in the SeZM NN descriptor by moving channel-only linear projections from grid resolution back to coefficient resolution (where possible), reducing work proportional to the grid size while preserving equivariant behavior.
Changes:
- Introduced a shared
_project_frames()helper to apply per-frameChannelLinearprojections directly on packed coefficient tensors. - Refactored
GridMLPandGridBranchto operate on coefficient operands and use injectedto_grid/from_gridprojectors only for the unavoidable point-wise grid product step. - Replaced the implicit GLU “identity” op path with an explicit
GridProductmodule and removed_apply_grid_op, unifying the grid-op call interface.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
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 (5)
✅ Files skipped from review due to trivial changes (1)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThis PR refactors grid-net computation in both PyTorch and DPModel implementations to operate at coefficient resolution. A new ChangesPyTorch Grid-Net Coefficient-Space Refactoring
DPModel Grid-Net Port and Grid Operation Classes
Documentation and Guard Clarifications
Test Parity and Coverage Updates
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5517 +/- ##
=======================================
Coverage 82.19% 82.19%
=======================================
Files 891 891
Lines 101599 101647 +48
Branches 4242 4240 -2
=======================================
+ Hits 83507 83552 +45
- Misses 16789 16792 +3
Partials 1303 1303 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
|
This PR passes its own CI but fails in the merge queue. The (run: https://github.com/deepmodeling/deepmd-kit/actions/runs/27460273106) Root cause — stale base + a semantic merge conflict. This PR makes deepmd-kit/source/tests/pt/model/test_dpa4_dpmodel_parity.py Lines 1629 to 1652 in 5d94bd6 Each side is fine alone, but the merge of this PR with current master (which already contains #5515) constructs the new required- Fix: rebase onto current master and update the now-stale call sites to pass |
![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
🤖 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/descriptor/dpa4_nn/grid_net.py`:
- Around line 692-695: The `op_type` parameter docstring in the GridNet class
(around lines 692-695) incorrectly states that `"mlp"` is not ported, but this
contradicts the actual implementation which supports all three operation types
including `"mlp"`. Update the docstring for the `op_type` parameter to remove
the note claiming `"mlp"` is not ported, and ensure it accurately reflects that
`"mlp"` is a supported option alongside `"glu"` and `"branch"`, consistent with
the module docstring, BaseGridNet implementation, and serialize/deserialize
logic.
🪄 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: 4227543e-5cfe-4584-b0ec-2e8d377dd576
📒 Files selected for processing (5)
deepmd/dpmodel/descriptor/dpa4_nn/block.pydeepmd/dpmodel/descriptor/dpa4_nn/ffn.pydeepmd/dpmodel/descriptor/dpa4_nn/grid_net.pydeepmd/pt/model/descriptor/sezm_nn/grid_net.pysource/tests/pt/model/test_dpa4_dpmodel_parity.py
💤 Files with no reviewable changes (1)
- source/tests/pt/model/test_dpa4_dpmodel_parity.py
✅ Files skipped from review due to trivial changes (2)
- deepmd/dpmodel/descriptor/dpa4_nn/block.py
- deepmd/dpmodel/descriptor/dpa4_nn/ffn.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/descriptor/sezm_nn/grid_net.py
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 1
🤖 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/descriptor/dpa4_nn/grid_net.py`:
- Around line 692-695: The `op_type` parameter docstring in the GridNet class
(around lines 692-695) incorrectly states that `"mlp"` is not ported, but this
contradicts the actual implementation which supports all three operation types
including `"mlp"`. Update the docstring for the `op_type` parameter to remove
the note claiming `"mlp"` is not ported, and ensure it accurately reflects that
`"mlp"` is a supported option alongside `"glu"` and `"branch"`, consistent with
the module docstring, BaseGridNet implementation, and serialize/deserialize
logic.
🪄 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: 4227543e-5cfe-4584-b0ec-2e8d377dd576
📒 Files selected for processing (5)
deepmd/dpmodel/descriptor/dpa4_nn/block.pydeepmd/dpmodel/descriptor/dpa4_nn/ffn.pydeepmd/dpmodel/descriptor/dpa4_nn/grid_net.pydeepmd/pt/model/descriptor/sezm_nn/grid_net.pysource/tests/pt/model/test_dpa4_dpmodel_parity.py
💤 Files with no reviewable changes (1)
- source/tests/pt/model/test_dpa4_dpmodel_parity.py
✅ Files skipped from review due to trivial changes (2)
- deepmd/dpmodel/descriptor/dpa4_nn/block.py
- deepmd/dpmodel/descriptor/dpa4_nn/ffn.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/descriptor/sezm_nn/grid_net.py
🛑 Comments failed to post (1)
deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py (1)
692-695:
⚠️ Potential issue | 🟡 Minor | ⚡ Quick winStale docstring:
"mlp"is now ported.The docstring claims
"mlp"is not ported, but this contradicts the module docstring (lines 15-17), theBaseGridNetimplementation (lines 539-546), and the serialize/deserialize logic (lines 777, 844) which all supportop_type="mlp".📝 Proposed fix
op_type : str - Point-wise grid operation; ``"glu"`` or ``"branch"`` (``"mlp"`` is - not ported). + Point-wise grid operation: ``"glu"``, ``"mlp"``, or ``"branch"``.🤖 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/descriptor/dpa4_nn/grid_net.py` around lines 692 - 695, The `op_type` parameter docstring in the GridNet class (around lines 692-695) incorrectly states that `"mlp"` is not ported, but this contradicts the actual implementation which supports all three operation types including `"mlp"`. Update the docstring for the `op_type` parameter to remove the note claiming `"mlp"` is not ported, and ensure it accurately reflects that `"mlp"` is a supported option alongside `"glu"` and `"branch"`, consistent with the module docstring, BaseGridNet implementation, and serialize/deserialize logic.
Summary by CodeRabbit
Release Notes
Refactor
to_grid/from_gridcallables and generalized grid-op handling for quadratic/product, polynomial MLP, and branch routing.mlpandbranchoperations.Tests
mlp, with updated forward wiring.Documentation