[WIP] pd: support dpa3 dynamic shape for pd backend #4828
Conversation
HydrogenSulfate
force-pushed
the
support_pd_dynamic_shape
branch
from
1a5fe62 to
dd70e99
Compare
📝 WalkthroughWalkthroughThis change standardizes tensor dimension handling and indexing conventions across several modules, primarily affecting how edge and angle indices are structured and used. It introduces explicit reshaping, modifies the shape and construction of index tensors, and clarifies broadcasting logic, without altering core algorithms or public interfaces. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant RepFlowLayer
participant Utils
User->>RepFlowLayer: forward(nlist, nlist_mask, ...)
RepFlowLayer->>Utils: get_graph_index(nlist, nlist_mask, ...)
Utils-->>RepFlowLayer: edge_index [2, n_edge], angle_index [3, n_angle]
RepFlowLayer->>RepFlowLayer: _cal_hg_dynamic(..., owner=edge_index[0], ...)
RepFlowLayer->>Utils: aggregate(data, owners, average, num_owner)
Utils-->>RepFlowLayer: aggregated tensor
RepFlowLayer-->>User: output
Possibly related PRs
Suggested reviewers
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. 🔧 Pylint (3.3.7)deepmd/pd/model/descriptor/repflow_layer.pyNo files to lint: exiting. deepmd/pd/model/network/utils.pyNo files to lint: exiting. deepmd/pd/model/descriptor/repflows.pyNo files to lint: exiting. ✨ Finishing Touches
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pd/model/descriptor/repflow_layer.py (1)
752-752: Remove unused variable assignment.The variable
nallis assigned but never used in this scope. Consider removing this assignment to clean up the code.- nall = node_ebd_ext.shape[1]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/pd/model/descriptor/repflow_layer.py(9 hunks)deepmd/pd/model/descriptor/repflows.py(2 hunks)deepmd/pd/model/network/utils.py(5 hunks)
🧰 Additional context used
🧠 Learnings (2)
deepmd/pd/model/descriptor/repflows.py (2)
Learnt from: njzjz
PR: deepmodeling/deepmd-kit#4160
File: deepmd/dpmodel/utils/env_mat.py:52-64
Timestamp: 2024-10-08T15:32:11.479Z
Learning: Negative indices in `nlist` are properly handled by masking later in the computation, so they do not cause issues in indexing operations.
Learnt from: njzjz
PR: deepmodeling/deepmd-kit#4160
File: deepmd/dpmodel/utils/env_mat.py:52-64
Timestamp: 2024-09-24T01:59:37.973Z
Learning: Negative indices in `nlist` are properly handled by masking later in the computation, so they do not cause issues in indexing operations.
deepmd/pd/model/network/utils.py (2)
Learnt from: njzjz
PR: deepmodeling/deepmd-kit#4160
File: deepmd/dpmodel/utils/env_mat.py:52-64
Timestamp: 2024-09-24T01:59:37.973Z
Learning: Negative indices in `nlist` are properly handled by masking later in the computation, so they do not cause issues in indexing operations.
Learnt from: njzjz
PR: deepmodeling/deepmd-kit#4160
File: deepmd/dpmodel/utils/env_mat.py:52-64
Timestamp: 2024-10-08T15:32:11.479Z
Learning: Negative indices in `nlist` are properly handled by masking later in the computation, so they do not cause issues in indexing operations.
🧬 Code Graph Analysis (1)
deepmd/pd/model/network/utils.py (1)
source/tests/consistent/descriptor/test_dpa3.py (1)
data(78-134)
🪛 Ruff (0.11.9)
deepmd/pd/model/descriptor/repflow_layer.py
752-752: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
🪛 Flake8 (7.2.0)
deepmd/pd/model/descriptor/repflow_layer.py
[error] 752-752: local variable 'nall' is assigned to but never used
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (21)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (13)
deepmd/pd/model/descriptor/repflows.py (2)
518-519: LGTM: Tensor shape initialization updated to match new indexing conventions.The initialization of
edge_indexandangle_indexhas been correctly updated from shapes[1, 3]to[2, 1]and[3, 1]respectively, which aligns with the new tensor dimension conventions where these tensors are structured as[2, n_edge]and[3, n_angle].
570-570: LGTM: Indexing updated to match new tensor shape convention.The change from
edge_index[:, 0](column-wise indexing) toedge_index[0](row-wise indexing) is consistent with the new tensor shape whereedge_indexhas dimensions[2, n_edge]instead of[n_edge, 2].deepmd/pd/model/network/utils.py (5)
32-43: LGTM: Optimized bin_count computation for better performance.This optimization computes
bin_countonly when needed (num_ownerisNoneor averaging is requested), which can improve performance in cases where bincount computation is expensive and unnecessary.
46-50: LGTM: Improved tensor initialization and assertion for safety.The output tensor initialization now consistently uses
num_ownerfor the first dimension, and the assertion ensuresbin_countis notNonebefore division when averaging is requested, preventing potential runtime errors.
59-59: LGTM: Added use_loc_mapping parameter for flexible index calculation.The new parameter provides control over how
frame_shiftis computed, allowing the function to work with different indexing schemes based on whether local mapping is used.
109-111: LGTM: Frame shift calculation adapted for different mapping modes.The conditional logic correctly adjusts the frame shift calculation based on
use_loc_mapping, using eithernallornlocas the multiplier, which maintains proper indexing behavior across different execution modes.
140-143: LGTM: Tensor stacking changes align with new indexing conventions.The change from concatenation to stacking transforms the output tensor shapes from
[n_edge, 2]and[n_angle, 3]to[2, n_edge]and[3, n_angle]respectively, which standardizes the indexing convention across the codebase.deepmd/pd/model/descriptor/repflow_layer.py (6)
375-377: LGTM: Improved code clarity with explicit tensor operations.The change from ellipsis-based indexing to explicit
.unsqueeze()calls makes the tensor operations more explicit and easier to understand, improving code readability.
588-592: LGTM: Standardized tensor reshaping with list arguments.Using list arguments for
reshape()instead of tuple unpacking is more explicit and consistent with modern tensor operation practices.
703-704: LGTM: Documentation updated to reflect new tensor shapes.The function signature documentation correctly reflects the new tensor shapes
2 x n_edgeand3 x n_angle, maintaining consistency with the implementation changes.
764-769: LGTM: Updated tensor indexing to match new shape conventions.The change from column-wise indexing (e.g.,
edge_index[:, 0]) to row-wise indexing (e.g.,edge_index[0]) correctly adapts to the new tensor shapes whereedge_indexhas dimensions[2, n_edge]andangle_indexhas dimensions[3, n_angle].
900-900: LGTM: Consistent use of explicit reshaping.The change to use list arguments in
reshape()maintains consistency with the other tensor operations in the codebase.
1046-1049: LGTM: Explicit tensor operations improve readability.The use of explicit
.unsqueeze()calls instead of implicit broadcasting makes the tensor operations more transparent and easier to debug.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4828 +/- ##
==========================================
- Coverage 84.57% 84.56% -0.01%
==========================================
Files 699 699
Lines 68036 68040 +4
Branches 3540 3540
==========================================
Hits 57540 57540
- Misses 9361 9364 +3
- Partials 1135 1136 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
HydrogenSulfate
changed the title
support running
input_torch_dynamic.jsonwith paddle backend(including CINN)TODO list:
masked_select_double_gradin eager mode PaddlePaddle/Paddle#73601masked_select_gradfor dynamic shape PaddlePaddle/Paddle#73622index_add_gradfor static decomposition PaddlePaddle/Paddle#73737index_put_gradfor static decomposition PaddlePaddle/Paddle#73747Summary by CodeRabbit
Bug Fixes
Refactor
New Features