feat(pt/dp/jax/xp): add DPA3 descriptor #4609
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
CodeQL found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.
| @staticmethod | ||
| def _cal_hg( |
There was a problem hiding this comment.
We can merge these three similar methods with dpa2 descriptor. For simplicity and staticmethod here, I kept them and may merge them in another PR.
| if not self.optim_update: | ||
| # nb x nloc x nnei x (n_dim * 2 + e_dim) | ||
| edge_info = torch.cat( | ||
| [ | ||
| torch.tile(node_ebd.unsqueeze(-2), [1, 1, self.nnei, 1]), | ||
| nei_node_ebd, | ||
| edge_ebd, | ||
| ], | ||
| dim=-1, | ||
| ) | ||
| else: | ||
| edge_info = None |
There was a problem hiding this comment.
I kept the original update here for clarity and future optimization test. The default method is within optim_update. We can remove this branch if necessary in the future.
| fix_stat_std : float, optional | ||
| If non-zero (default is 0.3), use this constant as the normalization standard deviation | ||
| instead of computing it from data statistics. | ||
| skip_stat : bool, optional | ||
| (Deprecated, kept only for compatibility.) This parameter is obsolete and will be removed. | ||
| If set to True, it forces fix_stat_std=0.3 for backward compatibility. | ||
| Transition to fix_stat_std parameter immediately. |
There was a problem hiding this comment.
Note the relationship between skip_stat and fix_stat_std. I kept skip_stat for compatibility with already-trained models. We can remove skip_stat in the future.
| n_multi_edge_message : int, optional | ||
| The head number of multiple edge messages to update node feature. | ||
| Default is 1, indicating one head edge message. |
There was a problem hiding this comment.
I kept n_multi_edge_message here also for compatibility, it exists in already-trained models with default value 1. I kept it with correct implemetation.
📝 WalkthroughWalkthroughThe pull request introduces a new descriptor called Changes
Sequence Diagram(s)sequenceDiagram
participant U as User/Input
participant D as DescrptDPA3 Descriptor
participant R as RepFlowLayer/Repflows
participant S as Serialization Engine
U->>D: Call with extended coordinates, atom types, and nlist
D->>R: Process message passing (node, edge, angle embeddings)
R-->>D: Return updated embeddings
D->>S: Serialize state / stats if needed
S-->>D: Return serialized data
D->>U: Return computed descriptor outputs (embeddings, rotations, stats)
Possibly related PRs
Suggested reviewers
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
⏰ Context from checks skipped due to timeout of 90000ms (20)
🔇 Additional comments (5)
✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
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: 14
🧹 Nitpick comments (13)
deepmd/jax/descriptor/repflows.py (1)
23-41: Consider clarifying the rationale behind double serialization.In the
__setattr__method ofDescrptBlockRepflows, line 31 re-serializes each layer withlayer.serialize()before deserializing again. While this may be intentional to ensure proper object structure, adding a short comment would help future maintainers understand the necessity of this procedure.deepmd/pt/model/descriptor/dpa3.py (3)
347-347: Rename unused loop variable for clarity.The loop variable
iiis never used. To clarify intent, rename it to_or_i.- for ii, descrpt in enumerate(descrpt_list): + for _, descrpt in enumerate(descrpt_list):🧰 Tools
🪛 Ruff (0.8.2)
347-347: Loop control variable
iinot used within loop bodyRename unused
iito_ii(B007)
423-423: Remove unused variable "env_mat".The variable
env_matis assigned but never used, which can cause confusion and clutter.- env_mat = repflow_variable.pop("env_mat")🧰 Tools
🪛 Ruff (0.8.2)
423-423: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
🪛 GitHub Check: CodeQL
[notice] 423-423: Unused local variable
Variable env_mat is not used.
479-479: Remove unused assignment to "nall".The local variable
nallis never referenced; removing it can improve clarity.- nall = extended_coord.view(nframes, -1).shape[1] // 3🧰 Tools
🪛 Ruff (0.8.2)
479-479: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
🪛 GitHub Check: CodeQL
[notice] 479-479: Unused local variable
Variable nall is not used.deepmd/pt/model/descriptor/repflows.py (2)
460-460: Rename unused loop variable "idx".Renaming
idxto_idxclarifies that the variable is not used within the loop.- for idx, ll in enumerate(self.layers): + for _idx, ll in enumerate(self.layers):🧰 Tools
🪛 Ruff (0.8.2)
460-460: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
570-574: Consider using a ternary operator for conciseness.Replacing the multi-line
if-elseblock with a ternary operator can make the code more concise without loss of clarity.- if callable(merged): - sampled = merged() - else: - sampled = merged + sampled = merged() if callable(merged) else merged🧰 Tools
🪛 Ruff (0.8.2)
570-574: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
deepmd/dpmodel/descriptor/repflows.py (1)
436-436: Rename unused loop variable to underscore.The variable
idxis not referenced in the loop. Renaming it to_or_idxclarifies its unused status.-for idx, ll in enumerate(self.layers): +for _idx, ll in enumerate(self.layers):🧰 Tools
🪛 Ruff (0.8.2)
436-436: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
source/tests/array_api_strict/descriptor/dpa3.py (1)
22-31: LGTM! Consider adding docstring for clarity.The implementation correctly handles attribute assignments with appropriate type conversions and deserialization. The class registration and inheritance are well-structured.
Consider adding a docstring to explain the purpose of the class and its attribute handling:
@BaseDescriptor.register("dpa3") class DescrptDPA3(DescrptDPA3DP): + """Array API Strict implementation of DPA3 descriptor. + + Handles attribute assignments with appropriate type conversions: + - mean, stddev: Converted to array API strict arrays + - repflows: Deserialized using DescrptBlockRepflows + - type_embedding: Deserialized using TypeEmbedNet + """ def __setattr__(self, name: str, value: Any) -> None:deepmd/jax/descriptor/dpa3.py (1)
23-35: LGTM! Consider adding docstring and return type hint.The implementation correctly handles attribute assignments with appropriate type conversions and deserialization. The class registration, inheritance, and decorators are well-structured.
Consider these improvements:
- Add a docstring to explain the purpose of the class and its attribute handling:
@BaseDescriptor.register("dpa3") @flax_module class DescrptDPA3(DescrptDPA3DP): + """JAX implementation of DPA3 descriptor. + + Handles attribute assignments with appropriate type conversions: + - mean, stddev: Converted to JAX arrays and wrapped in ArrayAPIVariable + - repflows: Deserialized using DescrptBlockRepflows + - type_embedding: Deserialized using TypeEmbedNet + """ def __setattr__(self, name: str, value: Any) -> None:
- Add return type hint for
__setattr__:- def __setattr__(self, name: str, value: Any) -> None: + def __setattr__(self, name: str, value: Any) -> None | Any:source/tests/consistent/descriptor/test_dpa3.py (2)
154-155: Remove commented code and simplify skip condition.The commented code suggests that the PD backend skip condition should be based on installation status and precision, but it's currently hardcoded to
True. Either implement the commented logic or simplify the code.- # return not INSTALLED_PD or precision == "bfloat16" - return True + return not INSTALLED_PD or precision == "bfloat16"
127-136: Remove unused parameters in skip_pt method.The method extracts parameters from
self.parambut doesn't use them. Consider removing the unused parameters to improve code clarity.- ( - update_residual_init, - exclude_types, - update_angle, - a_compress_rate, - a_compress_e_rate, - a_compress_use_split, - optim_update, - fix_stat_std, - n_multi_edge_message, - precision, - ) = self.param return CommonTest.skip_pt🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 127-127: Unused local variable
Variable update_residual_init is not used.
[notice] 128-128: Unused local variable
Variable exclude_types is not used.
[notice] 129-129: Unused local variable
Variable update_angle is not used.
[notice] 130-130: Unused local variable
Variable a_compress_rate is not used.
[notice] 131-131: Unused local variable
Variable a_compress_e_rate is not used.
[notice] 132-132: Unused local variable
Variable a_compress_use_split is not used.
[notice] 133-133: Unused local variable
Variable optim_update is not used.
[notice] 134-134: Unused local variable
Variable fix_stat_std is not used.
[notice] 135-135: Unused local variable
Variable n_multi_edge_message is not used.
[notice] 136-136: Unused local variable
Variable precision is not used.source/tests/universal/pt/model/test_model.py (1)
92-92: Fix typo in variable name.The variable name
defalut_des_paramcontains a typo and should bedefault_des_param.-defalut_des_param = [ +default_des_param = [deepmd/utils/argcheck.py (1)
1426-1587: LGTM! Comprehensive repflow argument configuration.The repflow block configuration is well-defined with:
- Clear parameter documentation
- Appropriate default values
- Proper type validation
- Sensible parameter grouping
Consider adding parameter range constraints in documentation.
For numeric parameters like dimensions and rates, it would be helpful to document any valid ranges or constraints in the docstrings.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (20)
deepmd/dpmodel/descriptor/__init__.py(2 hunks)deepmd/dpmodel/descriptor/dpa3.py(1 hunks)deepmd/dpmodel/descriptor/repflows.py(1 hunks)deepmd/dpmodel/utils/network.py(1 hunks)deepmd/jax/descriptor/__init__.py(2 hunks)deepmd/jax/descriptor/dpa3.py(1 hunks)deepmd/jax/descriptor/repflows.py(1 hunks)deepmd/pt/model/descriptor/__init__.py(2 hunks)deepmd/pt/model/descriptor/dpa3.py(1 hunks)deepmd/pt/model/descriptor/repflow_layer.py(1 hunks)deepmd/pt/model/descriptor/repflows.py(1 hunks)deepmd/pt/utils/utils.py(1 hunks)deepmd/utils/argcheck.py(1 hunks)source/tests/array_api_strict/descriptor/dpa3.py(1 hunks)source/tests/array_api_strict/descriptor/repflows.py(1 hunks)source/tests/consistent/descriptor/test_dpa3.py(1 hunks)source/tests/pt/model/test_dpa3.py(1 hunks)source/tests/universal/dpmodel/descriptor/test_descriptor.py(4 hunks)source/tests/universal/pt/descriptor/test_descriptor.py(3 hunks)source/tests/universal/pt/model/test_model.py(15 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/consistent/descriptor/test_dpa3.py
39-43: Use ternary operator DescrptDPA3PD = None if INSTALLED_PD else None instead of if-else-block
(SIM108)
source/tests/pt/model/test_dpa3.py
209-209: Local variable model is assigned to but never used
Remove assignment to unused variable model
(F841)
source/tests/universal/dpmodel/descriptor/test_descriptor.py
474-474: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pt/model/descriptor/repflow_layer.py
307-307: Local variable e_dim is assigned to but never used
Remove assignment to unused variable e_dim
(F841)
534-534: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
788-788: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
deepmd/dpmodel/descriptor/dpa3.py
222-222: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
587-587: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
deepmd/pt/model/descriptor/repflows.py
99-99: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
359-359: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
460-460: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
570-574: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
deepmd/dpmodel/descriptor/repflows.py
77-77: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
352-352: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
436-436: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
863-863: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
1125-1125: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
deepmd/pt/model/descriptor/dpa3.py
75-75: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
347-347: Loop control variable ii not used within loop body
Rename unused ii to _ii
(B007)
423-423: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
479-479: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
🪛 GitHub Check: CodeQL
source/tests/consistent/descriptor/test_dpa3.py
[notice] 127-127: Unused local variable
Variable update_residual_init is not used.
[notice] 128-128: Unused local variable
Variable exclude_types is not used.
[notice] 129-129: Unused local variable
Variable update_angle is not used.
[notice] 130-130: Unused local variable
Variable a_compress_rate is not used.
[notice] 131-131: Unused local variable
Variable a_compress_e_rate is not used.
[notice] 132-132: Unused local variable
Variable a_compress_use_split is not used.
[notice] 133-133: Unused local variable
Variable optim_update is not used.
[notice] 134-134: Unused local variable
Variable fix_stat_std is not used.
[notice] 135-135: Unused local variable
Variable n_multi_edge_message is not used.
[notice] 136-136: Unused local variable
Variable precision is not used.
[notice] 143-143: Unused local variable
Variable update_residual_init is not used.
[notice] 144-144: Unused local variable
Variable exclude_types is not used.
[notice] 145-145: Unused local variable
Variable update_angle is not used.
[notice] 146-146: Unused local variable
Variable a_compress_rate is not used.
[notice] 147-147: Unused local variable
Variable a_compress_e_rate is not used.
[notice] 148-148: Unused local variable
Variable a_compress_use_split is not used.
[notice] 149-149: Unused local variable
Variable optim_update is not used.
[notice] 150-150: Unused local variable
Variable fix_stat_std is not used.
[notice] 151-151: Unused local variable
Variable n_multi_edge_message is not used.
[notice] 152-152: Unused local variable
Variable precision is not used.
[notice] 160-160: Unused local variable
Variable update_residual_init is not used.
[notice] 161-161: Unused local variable
Variable exclude_types is not used.
deepmd/pt/model/descriptor/repflow_layer.py
[notice] 307-307: Unused local variable
Variable e_dim is not used.
[notice] 534-534: Unused local variable
Variable nall is not used.
[notice] 788-788: Unused local variable
Variable nitem is not used.
deepmd/dpmodel/descriptor/dpa3.py
[notice] 587-587: Unused local variable
Variable env_mat is not used.
deepmd/dpmodel/descriptor/repflows.py
[notice] 863-863: Unused local variable
Variable nall is not used.
[notice] 1125-1125: Unused local variable
Variable nitem is not used.
deepmd/pt/model/descriptor/dpa3.py
[notice] 423-423: Unused local variable
Variable env_mat is not used.
[notice] 479-479: Unused local variable
Variable nall is not used.
⏰ Context from checks skipped due to timeout of 90000ms (18)
- GitHub Check: build_docker (_cu11, 11)
- GitHub Check: build_docker (12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (true)
🔇 Additional comments (15)
deepmd/jax/descriptor/repflows.py (1)
43-62: Overall class structure appears well-implemented.The custom overrides for
__setattr__inRepFlowLayerconsistently handle deserialization of diverse layer attributes. This approach should help maintain clarity and consistency across different layer types.deepmd/pt/model/descriptor/dpa3.py (1)
66-573: Solid overall implementation and integration.Aside from the minor static analysis issues, the class comprehensively covers descriptor initialization, message passing, environment handling, and serialization. The
share_paramsmechanism and type-map updates appear robust and well-tested. Good job ensuring the descriptor is self-contained and accommodates future enhancements.🧰 Tools
🪛 Ruff (0.8.2)
75-75: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
347-347: Loop control variable
iinot used within loop bodyRename unused
iito_ii(B007)
423-423: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
479-479: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
🪛 GitHub Check: CodeQL
[notice] 423-423: Unused local variable
Variable env_mat is not used.
[notice] 479-479: Unused local variable
Variable nall is not used.deepmd/pt/model/descriptor/repflows.py (1)
1-604: Well-structured design with comprehensive method coverage.The new
DescrptBlockRepflowsclass features a clear initialization flow, separation of logic via helper methods, and integration with message passing layers. The forward method is robust, handling environment matrix processing, neighbor lists, and angle calculations. Overall, the code is neatly arranged and should be maintainable moving forward.🧰 Tools
🪛 Ruff (0.8.2)
99-99: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
359-359: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
460-460: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
570-574: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
deepmd/dpmodel/descriptor/__init__.py (1)
8-10: LGTM!The import statement and
__all__list update forDescrptDPA3follow the established pattern and maintain consistency with other descriptors.Also applies to: 36-36
deepmd/jax/descriptor/__init__.py (1)
8-10: LGTM!The import statement and
__all__list update forDescrptDPA3follow the established pattern and maintain consistency with other descriptors.Also applies to: 33-33
deepmd/pt/model/descriptor/__init__.py (1)
16-18: LGTM! Clean integration of the DPA3 descriptor.The import statement and
__all__list addition follow the established pattern and maintain alphabetical order.Also applies to: 55-55
source/tests/universal/pt/descriptor/test_descriptor.py (1)
7-7: LGTM! Comprehensive test coverage for DPA3 descriptor.The changes properly integrate DPA3 into the test framework, ensuring it's tested alongside other descriptors.
Also applies to: 24-24, 45-45
source/tests/array_api_strict/descriptor/repflows.py (2)
22-36: LGTM! Clean implementation of DescrptBlockRepflows.The class properly handles array conversion and deserialization of attributes, with clear handling of special cases.
39-56: LGTM! Comprehensive implementation of RepFlowLayer.The class handles all neural network layer attributes correctly, with proper deserialization and array conversion.
deepmd/pt/utils/utils.py (1)
42-43: LGTM! Clean addition of SiLU activation function.The implementation correctly uses PyTorch's
F.siluand follows the established pattern for activation functions.deepmd/dpmodel/utils/network.py (1)
321-328:⚠️ Potential issueFix incorrect SiLU implementation.
The current implementation of SiLU (Sigmoid Linear Unit) is incorrect. SiLU, also known as Swish, should be implemented as x * sigmoid(x).
elif activation_function == "silu": def fn(x): xp = array_api_compat.array_namespace(x) - # generated by GitHub Copilot - return x / (1 + xp.exp(-x)) + return x * (1 / (1 + xp.exp(-x))) # x * sigmoid(x) return fnLikely an incorrect or invalid review comment.
source/tests/universal/pt/model/test_model.py (3)
13-13: LGTM! Import statements for DPA3 descriptor added correctly.The new imports for
DescrptDPA3,DescriptorParamDPA3, andDescriptorParamDPA3Listare properly integrated with the existing imports.Also applies to: 59-60
99-99: LGTM! DPA3 descriptor param added to default parameters.The
DescriptorParamDPA3is correctly added to the default descriptor parameters list.
124-124: LGTM! DPA3 descriptor properly integrated into test parameterization.The
DescrptDPA3descriptor and its parameters are correctly added to the parameterized tests across all model test classes, maintaining consistency with the existing test structure.Also applies to: 139-139, 228-228, 243-243, 327-327, 338-338, 422-422, 433-433, 736-736, 747-747, 829-829, 839-839
deepmd/utils/argcheck.py (1)
1357-1422: LGTM! Well-structured descriptor registration with clear documentation.The DPA3 descriptor registration follows the established pattern and includes comprehensive argument definitions with appropriate types, defaults and documentation.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
deepmd/dpmodel/descriptor/repflows.py (2)
338-349: Implementation placeholder for statistics.Both
compute_input_statsandget_statsraiseNotImplementedError. If these are planned features, consider implementing them or remove them if they're no longer needed.Would you like me to propose an interim implementation or open a new issue to track this as a TODO?
436-436: Rename unused loop variable for clarity.The loop variable
idxis not used, so renaming it to_or removing it may improve clarity.- for idx, ll in enumerate(self.layers): + for _, ll in enumerate(self.layers): ...🧰 Tools
🪛 Ruff (0.8.2)
436-436: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/dpmodel/descriptor/repflows.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/dpmodel/descriptor/repflows.py
77-77: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
352-352: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
436-436: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
927-927: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
1189-1189: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test Python (1, 3.9)
🔇 Additional comments (5)
deepmd/dpmodel/descriptor/repflows.py (5)
77-77: Avoid mutable default argument in constructor.This constructor parameter
exclude_types: list[tuple[int, int]] = []uses a mutable default list, which can lead to unexpected behavior across invocations.Below is a suggested fix:
- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None, ) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
77-77: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
352-352: Avoid mutable default argument in method definition.Similar to the constructor, the parameter
exclude_types: list[tuple[int, int]] = []can cause shared state issues.Below is a suggested fix:
- def reinit_exclude( - self, - exclude_types: list[tuple[int, int]] = [], - ) -> None: + def reinit_exclude( + self, + exclude_types: Optional[list[tuple[int, int]]] = None, + ) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
352-352: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
927-927: Remove unused local variable.Variable
nallis assigned but never used after assignment.- nall = node_ebd_ext.shape[1] node_ebd = node_ebd_ext[:, :nloc, :]🧰 Tools
🪛 Ruff (0.8.2)
927-927: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
1189-1189: Remove unused local variable.Variable
nitemis assigned but never used.- nitem = len(update_list) uu = update_list[0]🧰 Tools
🪛 Ruff (0.8.2)
1189-1189: Local variable
nitemis assigned to but never usedRemove assignment to unused variable
nitem(F841)
1-1: Overall well-structured implementation.Aside from the minor issues noted above, the file is well-organized and follows a clean approach to implementing a descriptor block with multiple layers. Great job!
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (3)
source/tests/universal/dpmodel/descriptor/test_descriptor.py (1)
474-474:⚠️ Potential issueReplace mutable default argument.
Using mutable default arguments in Python can lead to unexpected behavior. Replace the empty list default argument with
Noneand initialize it within the function.-def DescriptorParamDPA3(..., exclude_types=[], ...): +def DescriptorParamDPA3(..., exclude_types=None, ...): input_dict = { "repflow": RepFlowArgs( **{ ... } ), - "exclude_types": exclude_types, + "exclude_types": [] if exclude_types is None else exclude_types, ... } return input_dict🧰 Tools
🪛 Ruff (0.8.2)
474-474: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/dpmodel/descriptor/repflows.py (2)
352-352: 🛠️ Refactor suggestionAvoid mutable default argument in
reinit_exclude.Similarly, using an empty list for
exclude_typesmay lead to shared references across calls. Replace withNoneand initialize locally to prevent subtle bugs.- def reinit_exclude( - self, - exclude_types: list[tuple[int, int]] = [], - ) -> None: + def reinit_exclude( + self, + exclude_types: Optional[list[tuple[int, int]]] = None, + ) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
352-352: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
77-77: 🛠️ Refactor suggestionAvoid mutable default argument for
exclude_types.Using
list[tuple[int,int]] = []as a default can inadvertently share state between instances. Consider adopting aNonedefault and initializing the list in the method body.- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None, … ) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
77-77: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
🧹 Nitpick comments (5)
deepmd/dpmodel/descriptor/repflows.py (5)
53-83: Consider splitting the constructor into smaller methods.This constructor accepts a large number of parameters, making it somewhat cumbersome to maintain and test. Refactoring smaller logical parts (e.g., parsing edge/angle configurations, setting up embeddings, etc.) into helper methods would improve readability and maintainability.
🧰 Tools
🪛 Ruff (0.8.2)
77-77: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
438-438: Rename the unused loop index.Since
idxis never referenced inside the loop, renaming it to_underscores its unused nature and improves clarity.-for idx, ll in enumerate(self.layers): +for _, ll in enumerate(self.layers): node_ebd_ext = xp_take_along_axis(node_ebd, mapping, axis=1) ...🧰 Tools
🪛 Ruff (0.8.2)
438-438: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
870-1169: Consider refactoring thecallmethod to enhance readability.This method handles a substantial portion of message-passing logic, neighbor/angle computations, and embeddings. Extracting smaller sub-methods (e.g. edge update, angle update, node update) can improve clarity and testability.
🧰 Tools
🪛 Ruff (0.8.2)
929-929: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
929-929: Remove the unused local variablenall.The variable
nallis never used after its assignment. Removing it avoids confusion and minor memory overhead.- nall = node_ebd_ext.shape[1] node_ebd = node_ebd_ext[:, :nloc, :]🧰 Tools
🪛 Ruff (0.8.2)
929-929: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
1191-1191: Remove the unused local variablenitem.
nitemis assigned but never referenced. Removing it will help eliminate potential confusion about code intentions.- nitem = len(update_list) uu = update_list[0]🧰 Tools
🪛 Ruff (0.8.2)
1191-1191: Local variable
nitemis assigned to but never usedRemove assignment to unused variable
nitem(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/dpmodel/descriptor/repflows.py(1 hunks)source/tests/universal/dpmodel/descriptor/test_descriptor.py(4 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/universal/dpmodel/descriptor/test_descriptor.py
474-474: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/dpmodel/descriptor/repflows.py
77-77: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
352-352: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
438-438: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
929-929: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
1191-1191: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test Python (1, 3.9)
🔇 Additional comments (4)
source/tests/universal/dpmodel/descriptor/test_descriptor.py (3)
10-10: LGTM!The imports are correctly organized and placed in the appropriate sections.
Also applies to: 21-23
498-500: Verify angle cutoff ratios.The angle cutoff values are set to half of the edge cutoffs (
rcut/2andrcut_smth/2), and the angle selection is set to a quarter of the edge selection (sel//4). Please verify these ratios are intentional and optimal for the DPA3 descriptor's performance.
529-546: LGTM!The parameterized test configurations are comprehensive and well-structured, covering various combinations of parameters. The descriptor is correctly integrated into the test class.
Also applies to: 606-606
deepmd/dpmodel/descriptor/repflows.py (1)
210-216: Great approach for defining embedding layers.Initializing
edge_embdandangle_embdwithNativeLayerprovides a clean and flexible way to handle embedding transformations, ensuring consistency in precision, seeding, and parameter initialization.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
deepmd/dpmodel/descriptor/repflows.py (5)
338-348: Add implementation or document why these methods raise NotImplementedError.Both
compute_input_statsandget_statsmethods raiseNotImplementedErrorwithout explanation. If these are intended to be abstract methods, consider:
- Adding a docstring explaining why they're not implemented
- Using
@abstractmethoddecorator if they must be implemented by subclassesWould you like me to help implement these methods or add appropriate documentation?
929-929: Remove unused variable.The variable
nallis assigned but never used.Apply this diff:
- nall = node_ebd_ext.shape[1] node_ebd = node_ebd_ext[:, :nloc, :]🧰 Tools
🪛 Ruff (0.8.2)
929-929: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
1191-1191: Remove unused variable.The variable
nitemis assigned but never used.Apply this diff:
- nitem = len(update_list) uu = update_list[0]🧰 Tools
🪛 Ruff (0.8.2)
1191-1191: Local variable
nitemis assigned to but never usedRemove assignment to unused variable
nitem(F841)
778-828: Document performance benefits of optimized methods.The optimized methods
optim_angle_updateandoptim_edge_updatewould benefit from docstrings explaining:
- How they optimize the computation
- Expected performance improvements
- Any trade-offs compared to non-optimized versions
Consider adding docstrings like:
def optim_angle_update(self, ...) -> np.ndarray: """Optimized version of angle update computation. This method improves performance by [explain optimization strategy]. Benchmarks show [X]% improvement over the non-optimized version. Note: Only used when optim_update=True. """Also applies to: 829-869
382-382: Improve assertion error messages.The assertions would be more helpful in debugging if they included descriptive error messages explaining the expected values.
Apply these diffs:
- assert list(atype_embd.shape) == [nframes, nloc, self.n_dim] + assert list(atype_embd.shape) == [nframes, nloc, self.n_dim], \ + f"Expected shape [{nframes}, {nloc}, {self.n_dim}], got {list(atype_embd.shape)}"- assert (a_dim * a_compress_e_rate) % (2 * a_compress_rate) == 0, ( - f"For a_compress_rate of {a_compress_rate}, a_dim*a_compress_e_rate must be divisible by {2 * a_compress_rate}. " - f"Currently, a_dim={a_dim} and a_compress_e_rate={a_compress_e_rate} is not valid." + assert (a_dim * a_compress_e_rate) % (2 * a_compress_rate) == 0, \ + f"Invalid compression parameters: a_dim ({a_dim}) * a_compress_e_rate ({a_compress_e_rate}) " \ + f"must be divisible by 2 * a_compress_rate ({2 * a_compress_rate})"Also applies to: 586-589
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/dpmodel/descriptor/repflows.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/dpmodel/descriptor/repflows.py
77-77: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
352-352: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
438-438: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
929-929: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
1191-1191: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (18)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (python)
- GitHub Check: Test C++ (true)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (c-cpp)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
deepmd/dpmodel/descriptor/repflows.py (3)
438-438: Rename the unused loop variable for clarity.The variable
idxis never referenced inside the loop. Consider renaming it to_or_idxto make it clear it’s unused.-for idx, ll in enumerate(self.layers): +for _idx, ll in enumerate(self.layers): node_ebd_ext = xp_take_along_axis(node_ebd, mapping, axis=1) node_ebd, edge_ebd, angle_ebd = ll.call(...)🧰 Tools
🪛 Ruff (0.8.2)
438-438: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
357-467: Consider splitting thecallmethod to improve maintainability.This method has numerous steps for handling neighbor lists, building masks, computing angles, and updating embeddings, making it fairly large. Splitting it into smaller functions (e.g., one function for edge-related logic, another for angle-related logic) would enhance clarity and testability.
🧰 Tools
🪛 Ruff (0.8.2)
438-438: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
877-920: Refactor thecallmethod for clarity.Similar to the descriptor’s
call, this method handles a complex workflow: extracting neighbor embeddings, performing angle computations, partial updates, etc. Decomposing it into smaller, reusable sub-routines (e.g., self-update, edge-update, angle-update) can make maintenance easier and reduce cognitive load.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/dpmodel/descriptor/repflows.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/dpmodel/descriptor/repflows.py
77-77: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
352-352: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
438-438: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
935-935: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
1197-1197: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
🔇 Additional comments (4)
deepmd/dpmodel/descriptor/repflows.py (4)
77-77: Replace the mutable default argument with None.Using a list as a default argument could expose you to subtle bugs due to Python’s handling of mutable default values. Consider initializing it to
Noneand setting the list within the constructor body.Apply this diff to address the warning:
- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None, ... ) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
77-77: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
352-352: Replace the mutable default argument with None.This method also uses a list as a default argument, which can cause unintended side effects. Initialize it to
Noneand handle the default inside the method body.Apply this diff:
- def reinit_exclude( - self, - exclude_types: list[tuple[int, int]] = [], - ) -> None: + def reinit_exclude( + self, + exclude_types: Optional[list[tuple[int, int]]] = None, + ) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
352-352: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
935-935: Remove unused local variable.The variable
nallis assigned but never used. Removing it avoids confusion and aligns with best practices.- nall = node_ebd_ext.shape[1] node_ebd = node_ebd_ext[:, :nloc, :]🧰 Tools
🪛 Ruff (0.8.2)
935-935: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
1197-1197: Remove unused local variable.The variable
nitemis assigned but never used. Deleting it simplifies the method.- nitem = len(update_list) uu = update_list[0]🧰 Tools
🪛 Ruff (0.8.2)
1197-1197: Local variable
nitemis assigned to but never usedRemove assignment to unused variable
nitem(F841)
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4609 +/- ##
==========================================
+ Coverage 84.58% 84.74% +0.16%
==========================================
Files 680 687 +7
Lines 64510 65996 +1486
Branches 3539 3538 -1
==========================================
+ Hits 54563 55929 +1366
- Misses 8807 8925 +118
- Partials 1140 1142 +2 ☔ View full report in Codecov by Sentry. |
| # env_mat doesn't store any value | ||
| pass | ||
| elif name == "emask": | ||
| value = PairExcludeMask(value.ntypes, value.exclude_types) |
There was a problem hiding this comment.
add else to handle all the rest cases
There was a problem hiding this comment.
I think adding a simple else: pass doesn't make sense...
There was a problem hiding this comment.
explicitly write pass is also helpful.
| elif name in {"repflows"}: | ||
| value = DescrptBlockRepflows.deserialize(value.serialize()) | ||
| elif name in {"type_embedding"}: | ||
| value = TypeEmbedNet.deserialize(value.serialize()) |
There was a problem hiding this comment.
add else to handle all the rest cases
| if value is not None: | ||
| value = NativeLayer.deserialize(value.serialize()) | ||
| elif name in {"n_residual", "e_residual", "a_residual"}: | ||
| value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value] |
There was a problem hiding this comment.
add else to handle all the rest cases.
|
|
||
| def get_rcut_smth(self) -> float: | ||
| """Returns the radius where the neighbor information starts to smoothly decay to 0.""" | ||
| return self.rcut_smth |
There was a problem hiding this comment.
UTs are not covering some of the methods of the dpa3 class, as the codecov reported.
| Returns | ||
| ------- | ||
| descriptor: torch.Tensor | ||
| the descriptor of shape nb x nloc x n_dim. | ||
| invariant single-atom representation. | ||
| g2: torch.Tensor | ||
| invariant pair-atom representation. | ||
| h2: torch.Tensor | ||
| equivariant pair-atom representation. | ||
| rot_mat: torch.Tensor | ||
| rotation matrix for equivariant fittings | ||
| sw: torch.Tensor | ||
| The switch function for decaying inverse distance. |
There was a problem hiding this comment.
The docstring for __init__ should be moved to the class itself. https://numpydoc.readthedocs.io/en/latest/format.html#class-docstring
| assert self.repflows.e_rcut >= self.repflows.a_rcut | ||
| assert self.repflows.e_sel >= self.repflows.a_sel |
There was a problem hiding this comment.
I suggest a clear error message to explain the error.
| elif activation_function == "silu": | ||
|
|
||
| def fn(x): | ||
| xp = array_api_compat.array_namespace(x) | ||
| # generated by GitHub Copilot | ||
| return x / (1 + xp.exp(-x)) |
There was a problem hiding this comment.
We have activation function consistent tests that read the list of supported activation functions from deepmd.common.VALID_ACTIVATION
deepmd-kit/source/tests/consistent/test_activation.py
Lines 7 to 9 in 9e0928a
NOTE: examples, custom ops and docs will be updated in other PRs.
Summary by CodeRabbit
New Features
DescrptDPA3descriptor for advanced molecular simulations, enhancing flexibility and performance.Tests
DescrptDPA3descriptor and repflow functionalities, ensuring consistency and reliability across different backends.DescrptDPA3to validate its functionality and performance across various configurations.