feat(jax): zbl #4301
feat(jax): zbl #4301
Conversation
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| # if xp.any(uu < 0): | ||
| # raise Exception("coord go beyond table lower boundary") |
Check notice
Code scanning / CodeQL
Commented-out code Note
|
Warning Rate limit exceeded@njzjz has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 16 minutes and 38 seconds before requesting another review. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. 📝 Walkthrough📝 Walkthrough📝 WalkthroughWalkthroughThis pull request introduces significant updates across various files in the Changes
Possibly related PRs
Suggested labels
Suggested reviewers
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: 12
🧹 Outside diff range and nitpick comments (6)
deepmd/jax/model/model.py (2)
60-96: Add docstring and improve input validationThe function lacks documentation and proper input validation for required parameters.
Consider adding a docstring and input validation:
def get_zbl_model(data: dict) -> DPZBLModel: + """Create a DPZBLModel combining deep potential and pair table models. + + Parameters + ---------- + data : dict + Configuration dictionary containing: + - descriptor: dict with model descriptor settings + - fitting_net: dict with fitting network settings + - use_srtab: str, path to the pair table file + - sw_rmin: float, minimum radius for switching function + - sw_rmax: float, maximum radius for switching function + - type_map: list[str], atom type mapping + - atom_exclude_types: list, optional + - pair_exclude_types: list, optional + + Returns + ------- + DPZBLModel + Combined model instance + + Raises + ------ + ValueError + If required parameters are missing or fitting type is unknown + """ + required_keys = {'descriptor', 'fitting_net', 'use_srtab', 'sw_rmin', 'sw_rmax', 'type_map'} + missing_keys = required_keys - data.keys() + if missing_keys: + raise ValueError(f"Missing required parameters: {missing_keys}")
75-83: Consider caching the pair table modelThe
PairTabAtomicModelis initialized with a file path and likely loads data from disk. Consider:
- Implementing caching mechanism for the loaded pair table data
- Adding memory usage warnings for large pair tables
- Documenting performance implications in the model's docstring
deepmd/dpmodel/model/make_model.py (1)
478-485: LGTM! Improved robustness of distance calculations.The changes to use squared distances instead of
normis a good optimization that:
- Prevents NaN errors from JAX during norm calculation
- Improves performance by avoiding unnecessary square root operations
- Maintains mathematical correctness by comparing squared distances
Consider adding a unit test that specifically verifies this behavior with JAX to ensure the changes prevent the NaN errors as intended.
deepmd/dpmodel/atomic_model/pairtab_atomic_model.py (2)
336-336: Use standard integer type in type annotationsAnnotating
nsplinewithnp.int64might cause issues with type checkers expecting standard Python types. Consider usingintfor type annotations to enhance code clarity and compatibility.Apply this diff to update the type annotation:
-def _extract_spline_coefficient( - i_type: np.ndarray, - j_type: np.ndarray, - idx: np.ndarray, - tab_data: np.ndarray, - nspline: np.int64, +) -> np.ndarray: +def _extract_spline_coefficient( + i_type: np.ndarray, + j_type: np.ndarray, + idx: np.ndarray, + tab_data: np.ndarray, + nspline: int, +) -> np.ndarray:
372-372: Ensure consistent type casting with array namespaceCasting to built-in
intmay lead to inconsistencies across different array backends. For consistency and to prevent potential issues, consider casting toxp.int64.Apply this diff to update the casting:
-clipped_indices = xp.clip(expanded_idx, 0, nspline - 1).astype(int) +clipped_indices = xp.clip(expanded_idx, 0, nspline - 1).astype(xp.int64)deepmd/dpmodel/atomic_model/linear_atomic_model.py (1)
73-82: Simplify 'mapping_list' initializationConsider initializing
mapping_listdirectly as an instance variable to improve code clarity:- mapping_list = [] common_type_map = set(type_map) self.type_map = type_map for tpmp in sub_model_type_maps: if not common_type_map.issubset(set(tpmp)): err_msg.append( f"type_map {tpmp} is not a subset of type_map {type_map}" ) mapping_list.append(self.remap_atype(tpmp, self.type_map)) - self.mapping_list = mapping_list + self.mapping_list = [] + for tpmp in sub_model_type_maps: + if not common_type_map.issubset(set(tpmp)): + err_msg.append( + f"type_map {tpmp} is not a subset of type_map {type_map}" + ) + self.mapping_list.append(self.remap_atype(tpmp, self.type_map))
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (12)
deepmd/dpmodel/atomic_model/linear_atomic_model.py(7 hunks)deepmd/dpmodel/atomic_model/pairtab_atomic_model.py(8 hunks)deepmd/dpmodel/descriptor/dpa1.py(2 hunks)deepmd/dpmodel/model/make_model.py(1 hunks)deepmd/dpmodel/utils/nlist.py(1 hunks)deepmd/dpmodel/utils/safe_gradient.py(1 hunks)deepmd/jax/atomic_model/linear_atomic_model.py(1 hunks)deepmd/jax/atomic_model/pairtab_atomic_model.py(1 hunks)deepmd/jax/model/__init__.py(1 hunks)deepmd/jax/model/dp_zbl_model.py(1 hunks)deepmd/jax/model/model.py(3 hunks)source/tests/consistent/model/test_zbl_ener.py(5 hunks)
🧰 Additional context used
🪛 GitHub Check: CodeQL
deepmd/dpmodel/atomic_model/pairtab_atomic_model.py
[notice] 281-282: Commented-out code
This comment appears to contain commented-out code.
🪛 Ruff
deepmd/dpmodel/utils/nlist.py
228-228: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/dpmodel/utils/safe_gradient.py
22-25: Use ternary operator mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis) instead of if-else-block
Replace if-else-block with mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis)
(SIM108)
🔇 Additional comments (22)
deepmd/jax/model/__init__.py (1)
2-4: LGTM! Clean module organization.
The changes properly expose the new DPZBLLinearEnergyAtomicModel while maintaining good Python packaging practices. The import structure is clean and the __all__ list is correctly updated to include the new model.
Also applies to: 9-12
deepmd/jax/atomic_model/pairtab_atomic_model.py (3)
1-16: LGTM: Well-organized imports with proper licensing.
The file has appropriate LGPL-3.0-or-later licensing and well-structured imports grouped by source.
19-20: LGTM: Proper Flax module definition.
The class correctly extends PairTabAtomicModelDP and is properly decorated with @flax_module for Flax compatibility.
1-27: Verify the implementation of mentioned methods.
The AI summary mentions several methods that are not visible in the provided code:
- forward_atomic
- _pair_tabulated_inter
- _extract_spline_coefficient
✅ Verification successful
Implementation follows expected inheritance pattern
The JAX implementation correctly inherits from PairTabAtomicModelDP which contains all the required methods:
forward_atomic_pair_tabulated_inter_extract_spline_coefficient
The JAX class only overrides __setattr__ to handle JAX-specific array conversions while inheriting the core functionality from the parent class, which is the expected pattern for framework-specific implementations.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the existence and implementation of the mentioned methods
# Expected: Find method implementations in the file
echo "Searching for method implementations..."
rg -U "def (forward_atomic|_pair_tabulated_inter|_extract_spline_coefficient)" "deepmd/jax/atomic_model/pairtab_atomic_model.py"
Length of output: 217
Script:
#!/bin/bash
# Let's check the parent class implementation and inheritance chain
echo "Checking parent class implementation..."
rg -U "class PairTabAtomicModelDP" "deepmd/dpmodel/atomic_model/pairtab_atomic_model.py" -A 10
echo -e "\nChecking for method implementations in parent class..."
rg -U "def (forward_atomic|_pair_tabulated_inter|_extract_spline_coefficient)" "deepmd/dpmodel/atomic_model/pairtab_atomic_model.py"
echo -e "\nChecking if methods are implemented in any base classes..."
fd "atomic_model" --type f | xargs rg -U "def (forward_atomic|_pair_tabulated_inter|_extract_spline_coefficient)"
Length of output: 1573
deepmd/dpmodel/utils/safe_gradient.py (3)
1-9: LGTM! Well-documented header with helpful reference.
The file header is properly licensed and documented, with a useful reference to JAX documentation explaining the gradient issues being addressed.
11-15: LGTM! Robust implementation of safe sqrt.
The implementation correctly handles edge cases and follows JAX's recommendations for safe gradients. Consider verifying the behavior with edge cases.
Let's verify the implementation with some test cases:
def test_safe_sqrt():
# Test cases:
# 1. x = 0.0 should return 0.0 with zero gradient
# 2. x < 0 should return 0.0
# 3. x > 0 should return sqrt(x)
import jax
import jax.numpy as jnp
x = jnp.array([4.0, 0.0, -1.0])
y = safe_for_sqrt(x)
# Expected: [2.0, 0.0, 0.0]
grad_fn = jax.grad(lambda x: safe_for_sqrt(x).sum())
grad = grad_fn(x)
# Expected: [0.25, 0.0, 0.0]18-32: LGTM! Consider a minor optimization.
The implementation correctly handles vector norm calculation with proper handling of edge cases and parameters.
Consider using a ternary operator for better readability:
- if keepdims:
- mask_squeezed = mask
- else:
- mask_squeezed = xp.squeeze(mask, axis=axis)
+ mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis)Let's verify the implementation with some test cases:
def test_safe_vector_norm():
# Test cases:
# 1. Zero vector should return 0 with zero gradient
# 2. Non-zero vector should return correct norm
# 3. Mixed case with axis and keepdims
import jax
import jax.numpy as jnp
x = jnp.array([[3.0, 4.0, 0.0], [0.0, 0.0, 0.0]])
y = safe_for_vector_norm(x, axis=1)
# Expected: [5.0, 0.0]
grad_fn = jax.grad(lambda x: safe_for_vector_norm(x, axis=1).sum())
grad = grad_fn(x)
# Expected: [[0.6, 0.8, 0.0], [0.0, 0.0, 0.0]]🧰 Tools
🪛 Ruff
22-25: Use ternary operator mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis) instead of if-else-block
Replace if-else-block with mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis)
(SIM108)
deepmd/jax/atomic_model/linear_atomic_model.py (3)
1-22: LGTM: Well-organized imports with proper licensing.
The imports are properly organized and the LGPL-3.0-or-later license is correctly declared.
25-26: LGTM: Proper class definition with Flax compatibility.
The class is correctly decorated with @flax_module for JAX/Flax compatibility and properly inherits from the base DP implementation.
1-38: Verify consistency with other JAX models.
Let's verify that this implementation follows the same patterns as other JAX models in the codebase.
✅ Verification successful
Let me try with corrected paths and more specific patterns.
Implementation is consistent with JAX model patterns
The implementation follows the established patterns in the JAX codebase:
- Uses the common
base_atomic_model_set_attrhelper like other JAX models - Consistently wraps JAX arrays with
ArrayAPIVariableandto_jax_arraysimilar toPairTabAtomicModel - Properly inherits and extends the DP model base class like other JAX atomic models
- Correctly handles model serialization/deserialization in the same way as other implementations
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check consistency of array API usage across JAX models
# Check array API usage patterns in other JAX models
echo "Checking array API usage patterns..."
rg -l "ArrayAPIVariable|to_jax_array" "deepmd/jax/**/*model*.py"
# Check base class method overrides
echo "Checking __setattr__ implementations..."
ast-grep --pattern 'class $_ {
$$$
__setattr__($_, $_) {
$$$
ArrayAPIVariable($$$)
$$$
}
$$$
}'
Length of output: 421
Script:
#!/bin/bash
# Check array API usage patterns in JAX models
echo "Checking array API usage patterns..."
rg -l "ArrayAPIVariable|to_jax_array" deepmd/jax/atomic_model/
# Check __setattr__ implementations in JAX models
echo "Checking __setattr__ implementations..."
rg -A 10 "__setattr__" deepmd/jax/atomic_model/
# Check base class implementations
echo "Checking base class patterns..."
rg -A 5 "class.*AtomicModel" deepmd/jax/atomic_model/
Length of output: 4270
deepmd/jax/model/model.py (2)
6-26: LGTM: Imports are well-organized and necessary
The new imports are properly structured and all are utilized in the implementation of the ZBL model functionality.
111-112: Verify condition ordering and model compatibility
The current implementation checks for spin models before ZBL models. Please verify:
- If this ordering is intentional
- If ZBL and spin models can coexist
- If additional validation is needed when both 'spin' and 'use_srtab' are present in the data
Let's check if there are any existing tests or implementations combining spin and ZBL:
✅ Verification successful
Based on the test files and configuration examples found, I can now provide a final response:
Condition ordering is correct, no compatibility issues found
The current implementation is correct because:
- The ordering of conditions is intentional - checking for spin models before ZBL models is appropriate as they are independent features.
- The test files show that ZBL models are configured independently through the
use_srtabparameter and its associated settings (smin_alpha,sw_rmin,sw_rmax). - There's no evidence in the test files of any conflicts or special handling needed when both spin and ZBL features are present, suggesting they operate independently.
- The test configurations (particularly in
source/tests/universal/common/cases/model/model.py) show ZBL and spin models are tested separately, indicating they are designed to be independent features.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for any mentions of both spin and ZBL/srtab in tests or implementations
rg -l "spin.*srtab|srtab.*spin"
Length of output: 33
Script:
#!/bin/bash
# Let's search for the get_model function implementation to understand the context better
ast-grep --pattern 'def get_model($$$)'
# Let's also search for get_zbl_model implementation
ast-grep --pattern 'def get_zbl_model($$$)'
# Search for any configuration or test files containing 'srtab' to understand usage
rg -l "srtab"
# Search for any configuration or test files containing both ZBL references
rg -l "zbl"
# Let's also check for any documentation about model compatibility
fd "README|CONTRIBUTING|docs" --exec cat {}
Length of output: 21462
Script:
#!/bin/bash
# Let's search for the implementation of get_model and get_zbl_model to understand the order and compatibility
ast-grep --pattern 'def get_model($$$)' deepmd/jax/model/model.py
# Let's also check the implementation of get_zbl_model
ast-grep --pattern 'def get_zbl_model($$$)' deepmd/jax/model/model.py
# Search for any test files that might show model compatibility
rg -A 10 "use_srtab" source/tests/
Length of output: 9035
source/tests/consistent/model/test_zbl_ener.py (3)
16-16: LGTM: Import of JAX installation flag.
The import follows the established pattern for backend support flags.
31-35: LGTM: JAX model imports.
The conditional import block follows the established pattern and properly handles the case when JAX is not installed.
95-95: LGTM: JAX model class property.
The property follows the established pattern for backend model class references.
deepmd/dpmodel/descriptor/dpa1.py (2)
30-32: LGTM: Safe gradient import added.
The import of safe_for_vector_norm is correctly placed and follows the project's import style.
Line range hint 949-955: LGTM: Enhanced numerical stability with safe vector normalization.
The use of safe_for_vector_norm improves robustness by safely handling potential division by zero cases during vector normalization. This is a good practice for numerical stability.
The subsequent division operation also includes a safeguard using xp.maximum(normed, xp.full_like(normed, 1e-12)), which provides additional protection against division by very small numbers.
deepmd/jax/model/dp_zbl_model.py (1)
31-50: 🛠️ Refactor suggestion
Avoid redundant method if base class implementation suffices
If DPZBLModelDP already provides a suitable forward_common_atomic method, overriding it here may be unnecessary. Redundant methods can clutter the codebase and cause maintenance challenges. Unless there is a specific reason to override, consider removing this method to rely on the base class's implementation.
To confirm, please check if the parent class DPZBLModelDP has an appropriate forward_common_atomic method. If it does, and no additional functionality is added here, you can safely remove this method.
deepmd/dpmodel/utils/nlist.py (2)
236-239: Efficient calculation avoids NaN errors in JAX
Great job replacing the distance calculation with squared distances to prevent NaN errors when using JAX. This maintains correctness while improving robustness.
224-225: Verify the usage of xp.concat
Please ensure that xp.concat is the correct function for concatenation in the array API compatible namespace being used. In some array libraries, the correct function might be xp.concatenate.
Run the following script to check if xp has the concat attribute:
✅ Verification successful
Let me search for the definition and usage of concat in the array API namespace.
Let me try a different approach to search for the array API implementation and usage of concatenation functions.
Based on the search results, I can now generate the final response:
Both concat and concatenate are used in the codebase
The codebase uses both xp.concat and xp.concatenate functions interchangeably. This is safe because array libraries typically support both as aliases for the same operation. For example, in NumPy and other array libraries, concat is a commonly provided alias for concatenate.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify if `xp` has the `concat` function.
# Test: Search for `concat` method in the array API namespace. Expect: Definition of `concat`.
rg --files-with-matches 'def concat' $(fd -e py)
Length of output: 26861
Script:
#!/bin/bash
# Search for array API concat definition and usage
rg -A 5 "def concat" deepmd/dpmodel/array_api.py
Length of output: 50
Script:
#!/bin/bash
# Search for both concat and concatenate in array API related files
rg -l "concat|concatenate" deepmd/dpmodel/array_api.py deepmd/backend/jax.py deepmd/backend/pytorch.py deepmd/backend/tensorflow.py
# Also search for the usage context of concat in the nlist.py file
rg -B 2 -A 2 "xp.concat" deepmd/dpmodel/utils/nlist.py
Length of output: 845
deepmd/dpmodel/atomic_model/pairtab_atomic_model.py (2)
Line range hint 200-226: Efficient integration of array API compatibility
The updates to use xp as the array namespace and the adoption of array API functions improve compatibility and flexibility across different array backends. This enhances the maintainability and scalability of the code.
326-326: Robust computation of pairwise distances
Utilizing safe_for_sqrt ensures numerical stability when computing pairwise distances by preventing issues arising from negative inputs due to floating-point inaccuracies. This enhances the reliability of distance calculations.
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Outside diff range and nitpick comments (1)
deepmd/dpmodel/utils/safe_gradient.py (1)
22-26: Simplify conditional logic using ternary operator.The if-else block can be simplified for better readability.
- if keepdims: - mask_squeezed = mask - else: - mask_squeezed = xp.squeeze(mask, axis=axis) + mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis)🧰 Tools
🪛 Ruff
22-25: Use ternary operator
mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis)instead ofif-else-blockReplace
if-else-block withmask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis)(SIM108)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (1)
deepmd/dpmodel/utils/safe_gradient.py(1 hunks)
🧰 Additional context used
🪛 Ruff
deepmd/dpmodel/utils/safe_gradient.py
22-25: Use ternary operator mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis) instead of if-else-block
Replace if-else-block with mask_squeezed = mask if keepdims else xp.squeeze(mask, axis=axis)
(SIM108)
🔇 Additional comments (1)
deepmd/dpmodel/utils/safe_gradient.py (1)
1-9: LGTM! Well-documented module with clear licensing.
The module header provides excellent context about gradient handling with a reference to JAX documentation.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4301 +/- ##
==========================================
+ Coverage 84.38% 84.40% +0.02%
==========================================
Files 563 567 +4
Lines 52810 52910 +100
Branches 3054 3057 +3
==========================================
+ Hits 44564 44659 +95
- Misses 7287 7292 +5
Partials 959 959 ☔ View full report in Codecov by Sentry. |
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (3)
deepmd/jax/atomic_model/dp_atomic_model.py(3 hunks)deepmd/jax/atomic_model/linear_atomic_model.py(1 hunks)deepmd/jax/atomic_model/pairtab_atomic_model.py(1 hunks)
🔇 Additional comments (8)
deepmd/jax/atomic_model/dp_atomic_model.py (3)
4-4: LGTM: Import additions are appropriate
The new imports are necessary for JAX integration and type hinting, following the project's conventions.
Also applies to: 17-20
37-53: LGTM: Well-structured method implementation
The method is well-typed and properly forwards to the parent implementation while handling the neighbor list gradient appropriately.
49-49: Verify consistent usage of stop_gradient
The use of stop_gradient on nlist is a significant choice that prevents gradient computation through the neighbor list. Let's verify this pattern is consistent across similar atomic models.
✅ Verification successful
Usage of stop_gradient on nlist is consistent across atomic models
The verification shows that stop_gradient is consistently applied to nlist in all atomic model implementations:
- deepmd/jax/atomic_model/dp_atomic_model.py
- deepmd/jax/atomic_model/linear_atomic_model.py
- deepmd/jax/atomic_model/pairtab_atomic_model.py
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for consistent usage of stop_gradient with nlist across atomic models
# Expected: All atomic models should apply stop_gradient to nlist in forward_common_atomic
# Search for forward_common_atomic methods that use nlist
ast-grep --pattern 'def forward_common_atomic($$$) {
$$$
nlist
$$$
}'
# Search for stop_gradient usage with nlist
rg -A 2 'stop_gradient.*nlist' deepmd/jax/atomic_model/
Length of output: 849
deepmd/jax/atomic_model/pairtab_atomic_model.py (3)
1-22: LGTM! Well-organized imports and proper licensing.
The file structure follows best practices with proper licensing and well-organized imports grouped by functionality.
24-25: LGTM! Proper class definition with Flax integration.
The class correctly inherits from the base DP model and uses the Flax module decorator for JAX compatibility.
34-50: Verify the necessity of stop_gradient for nlist parameter.
The implementation looks correct, but let's verify why only the nlist parameter needs gradient stopping while other parameters don't.
✅ Verification successful
Let me gather more context about the base class and its implementation to understand this pattern better.
Let me search for the base implementation and any documentation about the neighbor list handling.
Let me search for the neighbor list computation to understand why its gradient needs to be stopped.
Let me check the neighbor list related implementations and their base classes.
Based on the evidence gathered from the codebase, I can now provide a final response about the use of stop_gradient on nlist.
The use of stop_gradient on nlist parameter is correct and necessary.
The neighbor list (nlist) is a structural input that defines the spatial relationships between atoms. It's computed based on geometric criteria (like cutoff distances) and doesn't need gradient propagation during backpropagation because:
- It's a discrete structure that determines which atoms are considered neighbors
- The neighbor list computation is not a differentiable operation
- All atomic models (DPAtomicModel, PairTabAtomicModel, LinearAtomicModel) consistently apply this pattern
This is a common practice in molecular modeling where structural inputs that define connectivity or neighbor relationships are treated as static during gradient computation.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for other usages of stop_gradient with nlist to understand the pattern
rg -l "stop_gradient.*nlist" | while read -r file; do
echo "=== $file ==="
rg -A 2 -B 2 "stop_gradient.*nlist" "$file"
done
# Search for any comments or documentation explaining this behavior
rg -l -g "*.md" -g "*.rst" "nlist.*gradient" | while read -r file; do
echo "=== $file ==="
rg -A 2 -B 2 "nlist.*gradient" "$file"
done
Length of output: 1217
Script:
#!/bin/bash
# Find the base class implementation
ast-grep --pattern 'class $_AtomicModel'
# Look for any comments or docstrings in the atomic model files
rg -g "*.py" -B 5 "class.*AtomicModel" deepmd/jax/atomic_model/
# Check if there are any tests that might explain this behavior
fd "test.*atomic.*py" | while read -r file; do
echo "=== $file ==="
rg -A 2 -B 2 "forward_common_atomic|stop_gradient.*nlist" "$file"
done
Length of output: 11587
Script:
#!/bin/bash
# Look for base class implementation and any documentation about nlist
rg -g "*.py" -B 10 "class DPAtomicModelDP" deepmd/jax/atomic_model/
# Search for any comments or documentation about neighbor list handling
rg -g "*.py" "nlist" -C 5 deepmd/jax/atomic_model/
# Look for any gradient-related operations in the codebase
rg -g "*.py" "gradient|backward" -C 5 deepmd/jax/atomic_model/
Length of output: 6680
Script:
#!/bin/bash
# Search for neighbor list computation
rg -g "*.py" -B 10 "def compute.*nlist" deepmd/jax/
# Look for any documentation about neighbor list in docstrings
rg -g "*.py" -A 10 '""".*neighbor.*"""' deepmd/jax/
# Search for any imports or utilities related to neighbor list
rg -g "*.py" "import.*neighbor|class.*Neighbor" deepmd/jax/
Length of output: 1441
Script:
#!/bin/bash
# Look for the base neighbor list implementations
ast-grep --pattern 'class $_NeighborGatedAttention'
# Check the content of neighbor_stat.py
cat deepmd/jax/utils/neighbor_stat.py
# Look for any documentation in the descriptor module
rg -g "*.py" -B 10 "class NeighborGatedAttention" deepmd/jax/descriptor/
Length of output: 4309
deepmd/jax/atomic_model/linear_atomic_model.py (2)
32-43: Consider adding error handling for input validation and deserialization failures.
While the __setattr__ method assigns values based on attribute names, adding error handling would enhance robustness. Specifically:
- Validate input types for attributes like
mapping_list,zbl_weight, andmodels. - Handle potential exceptions during deserialization of the
models.
45-61: Verify the use of jax.lax.stop_gradient on nlist.
In the forward_common_atomic method, nlist is wrapped with jax.lax.stop_gradient, which stops gradients from flowing through nlist. Please confirm if stopping gradients on nlist is intentional, as it affects backpropagation and may impact model training.
|
are you planning to revert the modification on the nlist? |
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Reverted. |
Summary by CodeRabbit
Release Notes
New Features
DPZBLLinearEnergyAtomicModelandPairTabAtomicModel, enhancing atomic model functionalities.get_zbl_modelfunction for constructingDPZBLModelfrom input data.safe_for_vector_normandsafe_for_sqrt.Bug Fixes
format_nlistto prevent NaN errors.Documentation
Tests
test_zbl_ener.py.