feat(pt/dp): add dynamic sel for DPA3 #4754
Conversation
📝 Walkthrough## Walkthrough
This change introduces a "dynamic neighbor selection" feature to the DPA3 descriptor and associated repflow layers in both NumPy and PyTorch backends. It adds new configuration parameters (`use_dynamic_sel`, `sel_reduce_factor`), adapts core logic to support dynamic (unpadded) neighbor and angle selection, and implements new utility functions for graph indexing and aggregation. Corresponding tests and argument schemas are updated.
## Changes
| File(s) | Change Summary |
|---------------------------------------------------------------------------------------------|----------------|
| `deepmd/dpmodel/descriptor/dpa3.py`, `deepmd/pt/model/descriptor/dpa3.py`, `deepmd/pd/model/descriptor/dpa3.py` | Added `use_dynamic_sel` and `sel_reduce_factor` to `RepFlowArgs` and passed them to `DescrptBlockRepflows` in DPA3 descriptor classes. Serialization/deserialization and constructor logic updated accordingly. |
| `deepmd/dpmodel/descriptor/repflows.py`, `deepmd/pt/model/descriptor/repflows.py`, `deepmd/pd/model/descriptor/repflows.py` | Added dynamic selection support to `DescrptBlockRepflows` and `RepFlowLayer`. New parameters control neighbor selection and normalization. New dynamic methods and logic for flattened neighbor/angle handling, aggregation, and index management. In PD backend, dynamic selection is disallowed with an explicit error. |
| `deepmd/pt/model/descriptor/repflow_layer.py`, `deepmd/pd/model/descriptor/repflow_layer.py` | Added dynamic selection logic to `RepFlowLayer`, including new dynamic methods for rotation, symmetrization, and optimized updates. Modified `forward` (PyTorch) and serialization logic to support new parameters and dynamic paths. PD backend updated for serialization versioning and parameter storage. |
| `deepmd/dpmodel/utils/network.py`, `deepmd/pt/model/network/utils.py` | Added utility functions `aggregate` and `get_graph_index` for grouping/aggregation and index mapping in graph-based message passing. |
| `deepmd/dpmodel/utils/__init__.py` | Exported `aggregate` and `get_graph_index` from the utils module. |
| `deepmd/utils/argcheck.py` | Added `use_dynamic_sel` and `sel_reduce_factor` as optional arguments in the DPA3 repflow argument schema with documentation. |
| `source/tests/consistent/descriptor/test_dpa3.py`, `source/tests/universal/dpmodel/descriptor/test_descriptor.py` | Added `use_dynamic_sel` parameter to test parameterizations and logic. Adjusted test skipping and tolerance logic. |
| `source/tests/universal/dpmodel/model/test_model.py` | Integrated DPA3 descriptor and parameter classes into energy and spin model parameterized test sets. |
| `source/tests/pt/model/test_dynamic_sel.py` | Added new test class to verify consistency between static and dynamic neighbor selection in the DPA3 descriptor under various configurations. |
## Sequence Diagram(s)
```mermaid
sequenceDiagram
participant User
participant DPA3Descriptor
participant RepFlowArgs
participant RepFlowBlock
participant RepFlowLayer
participant Utils
User->>DPA3Descriptor: Initialize (with use_dynamic_sel, sel_reduce_factor)
DPA3Descriptor->>RepFlowArgs: Pass parameters
DPA3Descriptor->>RepFlowBlock: Pass use_dynamic_sel, sel_reduce_factor
RepFlowBlock->>RepFlowLayer: Pass use_dynamic_sel, sel_reduce_factor
User->>DPA3Descriptor: Compute descriptors
DPA3Descriptor->>RepFlowBlock: Forward (inputs)
RepFlowBlock->>Utils: get_graph_index (if dynamic)
RepFlowBlock->>RepFlowLayer: Forward (inputs, edge_index, angle_index)
RepFlowLayer->>Utils: aggregate (if dynamic)
RepFlowLayer->>RepFlowBlock: Return embeddings
RepFlowBlock->>DPA3Descriptor: Return descriptorsPossibly related PRs
Suggested reviewers
|
There was a problem hiding this comment.
Actionable comments posted: 7
🔭 Outside diff range comments (3)
deepmd/dpmodel/descriptor/repflows.py (2)
189-224: 🛠️ Refactor suggestionAdd parameter validation
Same as in the PyTorch backend: guard against invalid
sel_reduce_factorand incompatible flag combinations right inside__init__.
786-838: 🛠️ Refactor suggestionSanity-check dynamically-derived neighbour counts
self.dynamic_e_sel = self.nnei / self.sel_reduce_factor
self.dynamic_a_sel = self.a_sel / self.sel_reduce_factorIf
self.nneiora_selis not divisible bysel_reduce_factor, a silent float is produced but later used in integer contexts (**-0.5, divisions). Round or cast explicitly and document the intended behaviour.deepmd/pt/model/descriptor/repflow_layer.py (1)
710-775:⚠️ Potential issueMake
edge_index&angle_indextruly optional to preserve backward compatibilityThe docstring marks both tensors as optional, yet the function signature makes them mandatory.
All pre-existing call sites will now break unless they are patched simultaneously.- edge_index: torch.Tensor, # n_edge x 2 - angle_index: torch.Tensor, # n_angle x 3 + edge_index: Optional[torch.Tensor] = None, # n_edge x 2 + angle_index: Optional[torch.Tensor] = None, # n_angle x 3 )Immediately after the argument block:
- n_edge = int(nlist_mask.sum().item()) + if self.use_dynamic_sel: + assert edge_index is not None and angle_index is not None, ( + "edge_index and angle_index must be provided when " + "use_dynamic_sel=True" + ) + n_edge = int(nlist_mask.sum().item())This small change keeps the new functionality while preventing a breaking API change.
🧰 Tools
🪛 Ruff (0.11.9)
759-759: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
🧹 Nitpick comments (19)
deepmd/pt/model/descriptor/dpa3.py (1)
153-154: New dynamic selection parameters properly propagated to DescrptBlockRepflowsThe dynamic neighbor selection feature is correctly implemented by passing the parameters from the repflow_args to the constructor.
However, I recommend adding documentation for these parameters in the class docstring to help users understand their purpose and usage.
@@ -86,6 +86,10 @@ use_tebd_bias : bool, Optional Whether to use bias in the type embedding layer. type_map : list[str], Optional + A list of strings. Give the name to each type of atoms. + use_dynamic_sel : bool, Optional + Whether to use dynamic selection of neighbors. + sel_reduce_factor : float, Optional A list of strings. Give the name to each type of atoms.deepmd/dpmodel/descriptor/dpa3.py (2)
126-137: Doc-string section needs a quick pass for coherenceThe two new parameters are correctly documented, but they break the visual grouping of the previous bullet list – the “smooth_edge_update” paragraph is now split in half. A quick tidy-up (move the block to the end or re-indent) will keep the generated docs readable.
164-166: Add basic validation forsel_reduce_factor
sel_reduce_factoris silently stored even whenuse_dynamic_sel=False, and no guard prevents zero / negative values.
Consider enforcing a positive value only when the flag is enabled; in the off case a neutral1.0keeps later maths sane.+ if self.use_dynamic_sel: + if sel_reduce_factor <= 0: + raise ValueError("`sel_reduce_factor` must be positive.") + self.sel_reduce_factor = sel_reduce_factor + else: + self.sel_reduce_factor = 1.0Also applies to: 192-194
deepmd/pt/model/network/utils.py (1)
33-45: Potential divide-by-zero / wrong normalisation whennum_owner> actual owners
bin_count.where(bin_count != 0, 1)forces empty owner bins to1, so averaging leaves zeros unchanged – good.
However, whennum_owner>max(owners)+1, you append ones, meaning you will divide by 1 instead of 0 for those entirely-missing owners.
If a caller relies on zeros for missing owners, the result is silently wrong.Consider appending zeros and masking the division instead:
- bin_count = torch.cat([bin_count, bin_count.new_ones(difference)]) + bin_count = torch.cat([bin_count, bin_count.new_zeros(difference)]) … - if average: - output = (output.T / bin_count).T + if average: + # avoid NaN when bin_count == 0 + safe_bc = torch.where(bin_count == 0, torch.ones_like(bin_count), bin_count) + output = output / safe_bc.unsqueeze(-1)source/tests/pt/model/test_dynamic_sel.py (1)
31-33: Unused module-leveldtypevariable – can be removed
dtype = env.GLOBAL_PT_FLOAT_PRECISIONis never referenced after definition because the loop shadows it with anotherdtype.
Delete to avoid confusion.source/tests/consistent/descriptor/test_dpa3.py (1)
66-69: Parameterisation: single-value tuples could be plain literalsMinor: using
(True, )instead ofTrueinside@parameterizedexpands the test matrix unnecessarily (extra unit dimension).
It’s harmless but removing the tuple keeps the cartesian product size minimal.- (True, ), # a_compress_use_split + (True,), # or just True if no cartesian product intendeddeepmd/pt/model/descriptor/repflows.py (5)
200-236: Move validation logic closer to parameters
use_dynamic_sel,sel_reduce_factorand the newly-deriveddynamic_*_selare set, but no check protects users from passing nonsense values (e.g., negative factors or 0). A short input-sanity block would prevent NaNs later on.+ if sel_reduce_factor <= 0: + raise ValueError("`sel_reduce_factor` must be positive.")
445-448: Potentially clobbering a legitimate neighbour indexed “0”Padding indices
-1are rewritten to0.
If atom-id 0 really occurs as a neighbour, padded slots become indistinguishable from the genuine neighbour in tensors passed to downstream ops (even though the boolean mask still helps).
Consider replacing-1withnall(an out-of-range sentinel) or keeping-1and masking insideget_graph_index.
476-495: Minor optimisation: avoid reallocating large boolean tensorsInside the dynamic branch we build a four-dimensional mask
a_nlist_mask = ... & ...only to immediately use it once for indexing and then discard it.
Given typical neighbour counts, this tensor can be tens of MiBs.
Instead compute the 1-D boolean index viatorch.nonzeroor boolean-mask directly againstangle_inputto save memory.
495-498: PreferNoneover dummy tensors for unused indicesWhen dynamic selection is disabled you fabricate a dummy
edge_indexof shape(1,3).
PassingNoneconveys intent more clearly and removes two unnecessary gather operations executed inRepFlowLayer.call, with zero functional cost.
1267-1270: Remove unused local variablenall
nallis assigned but never referenced afterwards, triggering lint (ruff F841) and slightly confusing readers.- nall = node_ebd_ext.shape[1]deepmd/dpmodel/descriptor/repflows.py (4)
130-140: Replicate doc-string clarification from PyTorch backendFor consistency, also mention that
use_dynamic_sel=Truepresently requiressmooth_edge_update=True, otherwise aNotImplementedErrorwill be raised later in the layer code.
475-478: Possible collision of sentinel valueReplacing
-1with0may collide with real atom id 0. See analogous remark for the PyTorch file.
507-527: Memory footprint of temporary masksBuilding
a_nlist_mask[:, :, :, None] & a_nlist_mask[:, :, None, :]doubles memory for large systems. Consider boolean-masking the angle tensors in-place instead.
1267-1270: Unused variablenall— linter warningSame as in the PyTorch counterpart, delete the assignment.
- nall = node_ebd_ext.shape[1]🧰 Tools
🪛 Ruff (0.11.9)
1268-1268: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
deepmd/pt/model/descriptor/repflow_layer.py (4)
330-383: Useeinsumto avoid a large temporary tensor in_cal_hg_dynamic
flat_h2[..., None] * flat_edge_ebd[:, None, :]materialises an(n_edge, 3, e_dim)tensor just to reshape it again.
torch.einsumgenerates the same result while keeping memory overhead lower and improving readability.-# n_edge x 3 x e_dim -flat_h2g2 = (flat_h2[..., None] * flat_edge_ebd[:, None, :]).reshape( - -1, 3 * e_dim -) +# n_edge x 3 x e_dim ➜ reshaped to (n_edge, 3*e_dim) +flat_h2g2 = torch.einsum("ni,nj->nij", flat_h2, flat_edge_ebd).reshape( + -1, 3 * e_dim +)
759-760: Remove unused variablenall
nallis assigned but never referenced, triggering static-analysis warnings (ruff F841).
Simply drop the assignment:- nall = node_ebd_ext.shape[1]🧰 Tools
🪛 Ruff (0.11.9)
759-759: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
893-904: Letaggregate()compute the mean instead of manual divisionYou already know the effective neighbour count (
dynamic_e_sel).
Passingaverage=True(or implementing mean insideaggregate) avoids an extra division and guards against accidental mismatch betweendynamic_e_seland the real count.-aggregate( - node_edge_update, - n2e_index, - average=False, - num_owner=nb * nloc, -).reshape(nb, nloc, node_edge_update.shape[-1]) / self.dynamic_e_sel +aggregate( + node_edge_update, + n2e_index, + average=True, # mean inside the helper + num_owner=nb * nloc, +).reshape(nb, nloc, node_edge_update.shape[-1])
347-355: Spelling nit: “Flatted” → “Flattened”Minor docstring polish improves readability and future documentation extraction.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (13)
deepmd/dpmodel/descriptor/dpa3.py(5 hunks)deepmd/dpmodel/descriptor/repflows.py(27 hunks)deepmd/dpmodel/utils/__init__.py(2 hunks)deepmd/dpmodel/utils/network.py(1 hunks)deepmd/pt/model/descriptor/dpa3.py(1 hunks)deepmd/pt/model/descriptor/repflow_layer.py(19 hunks)deepmd/pt/model/descriptor/repflows.py(8 hunks)deepmd/pt/model/network/utils.py(1 hunks)deepmd/utils/argcheck.py(2 hunks)source/tests/consistent/descriptor/test_dpa3.py(10 hunks)source/tests/pt/model/test_dynamic_sel.py(1 hunks)source/tests/universal/dpmodel/descriptor/test_descriptor.py(3 hunks)source/tests/universal/dpmodel/model/test_model.py(6 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
deepmd/dpmodel/utils/__init__.py (2)
deepmd/dpmodel/utils/network.py (2)
aggregate(962-1000)get_graph_index(1003-1091)deepmd/pt/model/network/utils.py (2)
aggregate(10-45)get_graph_index(49-135)
source/tests/universal/dpmodel/model/test_model.py (3)
deepmd/dpmodel/descriptor/dpa3.py (1)
DescrptDPA3(235-659)deepmd/pt/model/descriptor/dpa3.py (1)
DescrptDPA3(65-567)source/tests/universal/dpmodel/descriptor/test_descriptor.py (1)
DescriptorParamDPA3(467-530)
deepmd/pt/model/network/utils.py (1)
deepmd/dpmodel/utils/network.py (2)
aggregate(962-1000)get_graph_index(1003-1091)
deepmd/pt/model/descriptor/repflows.py (2)
deepmd/pt/model/network/utils.py (1)
get_graph_index(49-135)deepmd/pt/model/descriptor/repflow_layer.py (3)
RepFlowLayer(38-1310)_cal_hg(288-328)_cal_hg_dynamic(331-383)
🪛 Ruff (0.11.9)
deepmd/dpmodel/descriptor/repflows.py
1268-1268: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pt/model/descriptor/repflow_layer.py
759-759: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (30)
- 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-win_amd64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (python)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (true)
🔇 Additional comments (12)
deepmd/dpmodel/utils/__init__.py (2)
18-19: New utility functions added for dynamic neighbor selectionThe addition of
aggregateandget_graph_indexfunctions from the network module makes them publicly available through the utils package.
58-59: Properly exposed new functions in all listGood practice - ensuring the newly imported functions are properly exposed in the
__all__list allows them to be imported directly from the package.source/tests/universal/dpmodel/descriptor/test_descriptor.py (3)
486-486: Added parameter for dynamic selectionThe new parameter
use_dynamic_selhas been properly added to the descriptor parameters.
511-511: Parameter correctly passed to RepFlowArgs constructorThe dynamic selection parameter is correctly passed to the underlying RepFlowArgs class.
542-543: Test configuration updated for new parametersThe test configuration has been updated to include the dynamic selection feature, properly testing both enabled and disabled states. The change to
a_compress_use_splitrestricts it to only test theTrueoption.Can you confirm the reason for removing the
Falseoption froma_compress_use_splittests? This change seems to reduce test coverage.Also applies to: 547-547
deepmd/utils/argcheck.py (1)
1500-1512: Well-designed dynamic selection feature with clear documentation.The new parameters
use_dynamic_selandsel_reduce_factorare well-documented with clear descriptions of their purpose and effects. The default values are sensible, with dynamic selection disabled by default for backward compatibility. The documentation effectively explains the benefit of dynamic neighbor selection - allowing users to set larger selection values safely while only considering actual neighbors within the cutoff radius.Also applies to: 1507-1512, 1613-1626
deepmd/dpmodel/utils/network.py (1)
962-1000:xp.bincount/xp.add.atare backend-specific – verify JAX & CuPy compatibilityThe helper relies on
xp.bincountandxp.add.at, neither of which are part of the 2022-12 Array-API spec.
NumPy works, but JAX, CuPy or a future tensorflow backend may not expose these names througharray_api_compat, causing a hard crash in dynamic-selection mode.Please double-check supported back-ends and, if needed, fall back to a pure-
xpscatter-add implementation (e.g.xp.zeros+ manual indexing) or conditional import.source/tests/universal/dpmodel/model/test_model.py (1)
7-13: Tests updated correctlyThe descriptor & parameter lists now cover
DescrptDPA3; parametrisation and imports look consistent with the existing pattern.Also applies to: 44-46, 98-100, 114-115, 187-188, 203-204
deepmd/pt/model/descriptor/repflows.py (3)
590-603: Check scale factor whensel_reduce_factor≠ neighbor countThe normalisation uses
scale_factor = (self.nnei / self.sel_reduce_factor)**(-0.5).If
sel_reduce_factoris merely a heuristic and the actual neighbour count differs frame-by-frame, the rotation matrix may be mis-scaled.
Consider computing the divisor from the per-owner edge count (available viaaggregate(..., average=True)), or at least add a doc-string note clarifying the approximation.
1600-1605:NotImplementedErrorwording & early-exitThis runtime check is great, but you can move it to the constructor (see first comment) and turn it into a
ValueError. Doing so fails fast during graph construction instead of midway through a forward pass.
1592-1597: Normalisation uses fixeddynamic_a_sel
aggregate(...)/(self.dynamic_a_sel**0.5)assumes every edge participates in exactlydynamic_a_selangles, which is only an approximation.
If the neighbour distribution is broad (e.g., surface atoms), this may bias updates. A per-edge reduction factor (angle_count_per_edge**-0.5) would be more accurate.deepmd/dpmodel/descriptor/repflows.py (1)
1590-1605: Early validation forsmooth_edge_updatewith dynamic selectionRun-time
NotImplementedErrorcan be promoted to an early constructor-level check.
| a_compress_e_rate, | ||
| a_compress_use_split, | ||
| optim_update, | ||
| use_dynamic_sel, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| a_compress_e_rate, | ||
| a_compress_use_split, | ||
| optim_update, | ||
| use_dynamic_sel, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| a_compress_e_rate, | ||
| a_compress_use_split, | ||
| optim_update, | ||
| use_dynamic_sel, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| a_compress_e_rate, | ||
| a_compress_use_split, | ||
| optim_update, | ||
| use_dynamic_sel, |
Check notice
Code scanning / CodeQL
Unused local variable Note test
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
deepmd/pt/model/descriptor/repflow_layer.py (1)
106-109:⚠️ Potential issueAdd validation for
sel_reduce_factorto prevent invalid operationsYour code computes
dynamic_e_selanddynamic_a_selby dividing bysel_reduce_factor, but there's no validation to ensure this value is positive and non-zero. This could lead to division-by-zero errors or negative values that would cause issues later when used as divisors or in negative exponents.self.use_dynamic_sel = use_dynamic_sel self.sel_reduce_factor = sel_reduce_factor +if self.sel_reduce_factor <= 0: + raise ValueError(f"`sel_reduce_factor` must be positive, got {self.sel_reduce_factor}") self.dynamic_e_sel = self.nnei / self.sel_reduce_factor self.dynamic_a_sel = self.a_sel / self.sel_reduce_factordeepmd/dpmodel/descriptor/repflows.py (1)
657-708: 🛠️ Refactor suggestionProtect against empty edge sets in
_cal_hg_dynamicThe
_cal_hg_dynamicfunction doesn't handle the case wheren_edge=0, which can occur with extreme exclusions. This could cause theaggregatefunction to fail with empty input.def _cal_hg_dynamic( flat_edge_ebd: np.ndarray, flat_h2: np.ndarray, flat_sw: np.ndarray, owner: np.ndarray, num_owner: int, nb: int, nloc: int, scale_factor: float, ) -> np.ndarray: xp = array_api_compat.array_namespace(flat_edge_ebd, flat_h2, flat_sw, owner) n_edge, e_dim = flat_edge_ebd.shape + + # Handle empty edge case + if n_edge == 0: + return xp.zeros((nb, nloc, 3, e_dim), dtype=flat_edge_ebd.dtype)
🧹 Nitpick comments (3)
deepmd/pt/model/descriptor/repflow_layer.py (2)
746-746: Remove unused variable assignmentThe variable
nallis assigned but never used in this scope.-nb, nloc, nnei = nlist.shape -nall = node_ebd_ext.shape[1] +nb, nloc, nnei = nlist.shape🧰 Tools
🪛 Ruff (0.11.9)
746-746: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
1049-1057: Consider handling empty edge sets in dynamic aggregationWhile this code correctly calculates edge angle updates with dynamic selection, it doesn't handle the edge case where there are no angles to process (
n_angle=0). This could happen in systems with extreme exclusions.# n_angle x e_dim weighted_edge_angle_update = edge_angle_update * a_sw.unsqueeze(-1) # n_edge x e_dim +if weighted_edge_angle_update.shape[0] == 0: + # Handle empty case + padding_edge_angle_update = torch.zeros((n_edge, self.e_dim), + dtype=edge_ebd.dtype, + device=edge_ebd.device) +else: padding_edge_angle_update = aggregate( weighted_edge_angle_update, eij2a_index, average=False, num_owner=n_edge, ) / (self.dynamic_a_sel**0.5)deepmd/dpmodel/descriptor/repflows.py (1)
1287-1287: Remove unused variable assignmentThe variable
nallis assigned but never used in this scope.-nb, nloc, nnei = nlist.shape -nall = node_ebd_ext.shape[1] +nb, nloc, nnei = nlist.shape🧰 Tools
🪛 Ruff (0.11.9)
1287-1287: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
deepmd/dpmodel/descriptor/repflows.py(27 hunks)deepmd/dpmodel/utils/__init__.py(2 hunks)deepmd/dpmodel/utils/network.py(1 hunks)deepmd/pt/model/descriptor/repflow_layer.py(19 hunks)source/tests/consistent/descriptor/test_dpa3.py(10 hunks)source/tests/pt/model/test_dynamic_sel.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- deepmd/dpmodel/utils/network.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/dpmodel/utils/init.py
- source/tests/pt/model/test_dynamic_sel.py
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/dpmodel/descriptor/repflows.py
1287-1287: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pt/model/descriptor/repflow_layer.py
746-746: 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] 137-137: Unused local variable
Variable use_dynamic_sel is not used.
[notice] 173-173: Unused local variable
Variable use_dynamic_sel is not used.
[notice] 190-190: Unused local variable
Variable use_dynamic_sel is not used.
[notice] 249-249: Unused local variable
Variable use_dynamic_sel is not used.
⏰ Context from checks skipped due to timeout of 90000ms (25)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- 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 C++ (false)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test C++ (true)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (cuda, cuda)
- 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: Analyze (c-cpp)
🔇 Additional comments (9)
source/tests/consistent/descriptor/test_dpa3.py (3)
66-68: Test parameters updated correctly to includeuse_dynamic_selThe parameterization has been updated to include the new dynamic selection feature, with appropriate constraints:
a_compress_use_splitis now fixed toTrueonlyuse_dynamic_selis added withTrueandFalseoptionsThis ensures proper test coverage of the new dynamic neighbor selection feature.
159-161: PD backend appropriately skipped for dynamic selectionThe skip condition correctly excludes tests with
use_dynamic_sel=Truefor the PD backend, properly documenting that this functionality isn't yet supported in that backend.
108-109: Newsmooth_edge_updateparameter added correctlyThe
smooth_edge_updateparameter is properly added to the repflow configuration with a value ofTrue, which is necessary for dynamic selection to work correctly, as seen in the implementation where dynamic selection requires smooth edge updates.deepmd/pt/model/descriptor/repflow_layer.py (3)
330-384: Implementation of dynamic neighbor computation looks correctThe new method
_cal_hg_dynamicprovides a parallel implementation that operates on flattened neighborhoods rather than padded tensors. The scale factor calculation and aggregation operations look properly implemented.One note is that for large molecular systems with many neighbors, this approach should significantly reduce memory usage compared to the padded version.
535-558: Fixed tensor split order in angle updateThe code now correctly splits the matrix weights and applies them to the corresponding embeddings. The previous implementation had an issue with the order of these operations.
879-891: Dynamic selection implementation handles aggregation properlyThe implementation correctly uses the
aggregateutility function to gather data from flattened tensors back into per-atom representations, applying proper scaling based ondynamic_e_sel. This is a key aspect of the dynamic selection approach.deepmd/dpmodel/descriptor/repflows.py (3)
128-138: Documentation comprehensively explains the new parametersThe documentation for
use_dynamic_selandsel_reduce_factoris very thorough and clearly explains:
- The purpose of dynamic selection (avoiding fixed padding)
- The impact on neighbor selection counts
- The relationship with normalization factors
This is excellent practice for API documentation.
837-841: Dynamic selection implementation reduces memory usage with flattened tensorsBy using flattened tensors for edge and angle embeddings with dynamic selection, this implementation avoids the memory overhead of padded tensors. For systems with many atoms but varying coordination numbers, this can significantly reduce memory usage while maintaining numerical accuracy.
1631-1635: Good error handling for required smooth edge updatesThe implementation correctly enforces that
smooth_edge_updatemust beTruewhen using dynamic selection by raising a clearNotImplementedErrorexception. This prevents users from using an unsupported configuration.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (4)
deepmd/pt/model/descriptor/repflow_layer.py (2)
106-109: Add missing validation forsel_reduce_factor
sel_reduce_factorneeds validation to prevent division by zero or negative values, which would cause errors when used in expressions likeself.dynamic_e_sel ** (-0.5).self.sel_reduce_factor = sel_reduce_factor +if self.sel_reduce_factor <= 0: + raise ValueError(f"`sel_reduce_factor` must be > 0, got {self.sel_reduce_factor}") self.dynamic_e_sel = self.nnei / self.sel_reduce_factor self.dynamic_a_sel = self.a_sel / self.sel_reduce_factor
1060-1064: Document dynamic selection requirementThe code requires
smooth_edge_update=Truewhen using dynamic selection, but this constraint isn't documented in user-facing documentation.Please update the user documentation to clearly state that when
use_dynamic_sel=True,smooth_edge_updatemust also beTrueto avoid the NotImplementedError.deepmd/dpmodel/descriptor/repflows.py (2)
837-840: Add missing validation forsel_reduce_factor
sel_reduce_factorneeds validation to prevent division by zero or negative values, which would cause errors when used in expressions likeself.dynamic_e_sel ** (-0.5).self.use_dynamic_sel = use_dynamic_sel self.sel_reduce_factor = sel_reduce_factor +if self.sel_reduce_factor <= 0: + raise ValueError(f"`sel_reduce_factor` must be > 0, got {self.sel_reduce_factor}") self.dynamic_e_sel = self.nnei / self.sel_reduce_factor self.dynamic_a_sel = self.a_sel / self.sel_reduce_factor
1631-1635: Document dynamic selection requirementThe code requires
smooth_edge_update=Truewhen using dynamic selection, but this constraint isn't documented in user-facing documentation.Please update the user documentation to clearly state that when
use_dynamic_sel=True,smooth_edge_updatemust also beTrueto avoid the NotImplementedError.
🧹 Nitpick comments (3)
deepmd/pt/model/descriptor/repflow_layer.py (1)
674-676: Fix shape reference inconsistencyThere's a small inconsistency in the tensor shape reference.
sub_node_ext_update = torch.index_select( - sub_node_ext_update.reshape(nf * nall, sub_node_update.shape[-1]), + sub_node_ext_update.reshape(nf * nall, sub_node_ext_update.shape[-1]), 0, n_ext2e_index, )While both shape references likely yield the same value, using the shape of the tensor being reshaped is more consistent and self-documenting.
deepmd/dpmodel/descriptor/repflows.py (2)
1202-1205: Fix shape reference inconsistencyThere's a small inconsistency in the tensor shape reference.
sub_node_ext_update = xp.take( - xp.reshape(sub_node_ext_update, (nf * nall, sub_node_update.shape[-1])), + xp.reshape(sub_node_ext_update, (nf * nall, sub_node_ext_update.shape[-1])), n_ext2e_index, axis=0, )While both shape references likely yield the same value, using the shape of the tensor being reshaped is more consistent and self-documenting.
1286-1288: Remove redundant computation ofnallThe variable
nallis computed twice, creating potential for divergence.nb, nloc, nnei = nlist.shape -nall = node_ebd_ext.shape[1] +# nall is already available from node_ebd_ext.shape[1] n_edge = int(nlist_mask.sum().item())The
nallis computed earlier at line 438 differently asnall = xp.reshape(coord_ext, (nframes, -1)).shape[1] // 3, but then assigned again here. Since you're already usingnode_ebd_ext.shape[1]elsewhere (like in line 1183), consider using just one consistent way of accessing this value throughout the code.🧰 Tools
🪛 Ruff (0.11.9)
1287-1287: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/dpmodel/descriptor/repflows.py(29 hunks)deepmd/pt/model/descriptor/repflow_layer.py(21 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
deepmd/dpmodel/descriptor/repflows.py (7)
deepmd/dpmodel/utils/network.py (2)
aggregate(962-1000)get_graph_index(1003-1097)deepmd/pt/model/network/utils.py (2)
aggregate(10-45)get_graph_index(49-135)deepmd/pt/model/descriptor/repflow_layer.py (6)
_cal_grrg(386-410)_cal_hg(288-328)_cal_hg_dynamic(331-383)symmetrization_op_dynamic(457-512)optim_angle_update_dynamic(556-604)optim_edge_update_dynamic(641-683)deepmd/dpmodel/descriptor/repformers.py (3)
_cal_grrg(750-775)_cal_hg(684-747)_make_nei_g1(615-645)deepmd/utils/version.py (1)
check_version_compatibility(2-27)source/tests/consistent/descriptor/test_dpa3.py (1)
data(75-125)source/tests/consistent/descriptor/test_se_e2_a.py (2)
data(77-97)data(321-341)
🪛 Ruff (0.11.9)
deepmd/dpmodel/descriptor/repflows.py
1287-1287: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pt/model/descriptor/repflow_layer.py
746-746: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (30)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- 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: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_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: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
🔇 Additional comments (1)
deepmd/dpmodel/descriptor/repflows.py (1)
695-708:Details
❌ Incorrect review comment
Handle empty edge case in dynamic calculation
The current implementation doesn't guard against the case when
n_edge==0, which is rare but possible with extreme exclusions.xp = array_api_compat.array_namespace(flat_edge_ebd, flat_h2, flat_sw, owner) n_edge, e_dim = flat_edge_ebd.shape +# Early return for empty edge case to avoid issues in aggregate function +if n_edge == 0: + return xp.zeros((nb, nloc, 3, e_dim), dtype=flat_edge_ebd.dtype) # n_edge x e_dim flat_edge_ebd = flat_edge_ebd * xp.expand_dims(flat_sw, axis=-1)Run the following to verify if there are tests covering the empty edge case:
🏁 Script executed:
#!/bin/bash # Check for any tests that might test the dynamic selection with empty edges grep -r "test.*dynamic.*sel" source/tests/ grep -r "empty.*edge" source/tests/ grep -r "n_edge.*==.*0" source/tests/ grep -r "no.*neighbor" source/tests/Length of output: 161
🏁 Script executed:
#!/bin/bash # Locate the definition of `aggregate` in the descriptor module grep -R "def aggregate" -n deepmd/dpmodel/descriptor # If it's in common.py, show the first 200 lines to inspect handling of empty input sed -n '1,200p' deepmd/dpmodel/descriptor/common.pyLength of output: 185
🏁 Script executed:
#!/bin/bash # List descriptor directory contents to locate files ls -R deepmd/dpmodel/descriptor # Search for all references to `aggregate` across the repo to find its definition or import grep -R "aggregate" -n . # Show the top of repflows.py to see how `aggregate` is imported sed -n '1,100p' deepmd/dpmodel/descriptor/repflows.pyLength of output: 12303
🏁 Script executed:
#!/bin/bash # Locate and display the aggregate function definition to check empty-input handling grep -n "def aggregate" -n deepmd/dpmodel/utils/network.py sed -n '950,1040p' deepmd/dpmodel/utils/network.pyLength of output: 2924
🏁 Script executed:
#!/bin/bash # Locate where `owner` and `num_owner` are defined in repflows.py grep -n "owner\s*=" -n deepmd/dpmodel/descriptor/repflows.py grep -n "num_owner\s*=" -n deepmd/dpmodel/descriptor/repflows.py # Show the block defining these variables (around line 650) sed -n '630,700p' deepmd/dpmodel/descriptor/repflows.pyLength of output: 3499
No need for an empty-edge guard in
_cal_hg_dynamic
Theaggregateimplementation indeepmd/dpmodel/utils/network.pyalready handlesn_edge == 0when you passnum_owner—it will pad the bin count tonum_ownerand return a zero tensor of shape(num_owner, feature_dim). That output reshapes naturally to(nb, nloc, 3, e_dim)and scales to zero without error. No additional early‐return is required here.Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (3)
deepmd/pt/model/network/utils.py (2)
96-107:⚠️ Potential issueAdd explicit boolean casting to avoid potential indexing errors
The function currently assumes
nlist_maskis a boolean tensor when used for indexing, but this isn't verified, which could lead to runtime errors if a different dtype is provided.Apply this fix to ensure robust indexing:
- n2e_index = n2e_index[nlist_mask] # graph node index, atom_graph[:, 0] + mask = nlist_mask.bool() + n2e_index = n2e_index[mask] # graph node index, atom_graph[:, 0] - n_ext2e_index = shifted_nlist[nlist_mask] # graph neighbor index, atom_graph[:, 1] + n_ext2e_index = shifted_nlist[mask] # graph neighbor index, atom_graph[:, 1]Similar changes should be applied to other boolean indexing operations at lines 112, 118, 123, and 128.
117-119:⚠️ Potential issueAdd explicit boolean casting for edge_index assignment
Similar to the previous issue, the boolean indexing for assignment at line 118 should use explicit boolean casting to avoid dtype-related errors.
Apply this fix:
- edge_index[nlist_mask] = edge_id + edge_index[mask] = edge_id # Reuse the mask variable created earlierdeepmd/dpmodel/descriptor/repflows.py (1)
666-718:⚠️ Potential issueMissing guard against empty edge sets in
_cal_hg_dynamic.When
n_edge==0(which can happen with extreme exclusions), the function will attempt to process an empty tensor which might cause issues in the underlying implementation.Add an early return statement to handle the case when no edges are found:
def _cal_hg_dynamic( flat_edge_ebd: np.ndarray, flat_h2: np.ndarray, flat_sw: np.ndarray, owner: np.ndarray, num_owner: int, nb: int, nloc: int, scale_factor: float, ) -> np.ndarray: xp = array_api_compat.array_namespace(flat_edge_ebd, flat_h2, flat_sw, owner) n_edge, e_dim = flat_edge_ebd.shape + # Return early if no edges are found + if n_edge == 0: + return xp.zeros((nb, nloc, 3, e_dim)) # n_edge x e_dim flat_edge_ebd = flat_edge_ebd * xp.expand_dims(flat_sw, axis=-1)
🧹 Nitpick comments (2)
deepmd/pt/model/descriptor/repflows.py (1)
484-504: Consider extracting common graph index logic to a separate methodThe dynamic selection branch adds complexity to the forward method. For better maintainability, consider extracting the dynamic selection logic to a separate helper method.
+ def _prepare_dynamic_selection( + self, + nlist, + nlist_mask, + a_nlist_mask, + nall, + edge_input, + h2, + sw, + angle_input + ): + # get graph index + edge_index, angle_index = get_graph_index( + nlist, nlist_mask, a_nlist_mask, nall + ) + # flat all the tensors + # n_edge x 1 + edge_input = edge_input[nlist_mask] + # n_edge x 3 + h2 = h2[nlist_mask] + # n_edge x 1 + sw = sw[nlist_mask] + # nb x nloc x a_nnei x a_nnei + a_nlist_mask_3d = a_nlist_mask[:, :, :, None] & a_nlist_mask[:, :, None, :] + # n_angle x 1 + angle_input = angle_input[a_nlist_mask_3d] + # n_angle x 1 + a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask_3d] + + return edge_index, angle_index, edge_input, h2, sw, angle_input, a_sw, a_nlist_mask_3dThen in the forward method:
if self.use_dynamic_sel: - # get graph index - edge_index, angle_index = get_graph_index( - nlist, nlist_mask, a_nlist_mask, nall - ) - # flat all the tensors - # n_edge x 1 - edge_input = edge_input[nlist_mask] - # n_edge x 3 - h2 = h2[nlist_mask] - # n_edge x 1 - sw = sw[nlist_mask] - # nb x nloc x a_nnei x a_nnei - a_nlist_mask = a_nlist_mask[:, :, :, None] & a_nlist_mask[:, :, None, :] - # n_angle x 1 - angle_input = angle_input[a_nlist_mask] - # n_angle x 1 - a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask] + edge_index, angle_index, edge_input, h2, sw, angle_input, a_sw, a_nlist_mask_3d = self._prepare_dynamic_selection( + nlist, nlist_mask, a_nlist_mask, nall, edge_input, h2, sw, angle_input, a_sw + ) + a_nlist_mask = a_nlist_mask_3ddeepmd/dpmodel/descriptor/repflows.py (1)
1296-1296: Unused variablenalldetected.The local variable
nallis assigned on line 1296 but never used in this scope.- nb, nloc, nnei = nlist.shape - nall = xp.reshape(coord_ext, (nframes, -1)).shape[1] // 3 + nb, nloc, nnei = nlist.shape🧰 Tools
🪛 Ruff (0.11.9)
1296-1296: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (10)
deepmd/dpmodel/descriptor/dpa3.py(5 hunks)deepmd/dpmodel/descriptor/repflows.py(29 hunks)deepmd/dpmodel/utils/network.py(1 hunks)deepmd/pd/model/descriptor/dpa3.py(1 hunks)deepmd/pd/model/descriptor/repflow_layer.py(5 hunks)deepmd/pd/model/descriptor/repflows.py(3 hunks)deepmd/pt/model/descriptor/repflows.py(8 hunks)deepmd/pt/model/network/utils.py(1 hunks)deepmd/utils/argcheck.py(2 hunks)source/tests/universal/dpmodel/descriptor/test_descriptor.py(3 hunks)
✅ Files skipped from review due to trivial changes (1)
- deepmd/dpmodel/utils/network.py
🚧 Files skipped from review as they are similar to previous changes (3)
- source/tests/universal/dpmodel/descriptor/test_descriptor.py
- deepmd/dpmodel/descriptor/dpa3.py
- deepmd/utils/argcheck.py
🧰 Additional context used
🧬 Code Graph Analysis (2)
deepmd/pd/model/descriptor/repflow_layer.py (2)
deepmd/utils/version.py (1)
check_version_compatibility(2-27)source/tests/consistent/descriptor/test_dpa3.py (1)
data(75-125)
deepmd/pt/model/network/utils.py (1)
deepmd/dpmodel/utils/network.py (2)
aggregate(962-1000)get_graph_index(1003-1097)
🪛 Ruff (0.11.9)
deepmd/dpmodel/descriptor/repflows.py
1296-1296: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (28)
- 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 cp310-manylinux_aarch64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test C++ (false)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (true)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
🔇 Additional comments (28)
deepmd/pt/model/network/utils.py (2)
9-45: Efficient implementation of row aggregation by owner indicesThe
aggregatefunction provides an efficient implementation for aggregating rows of a tensor by owner indices, supporting both summation and averaging modes. The code correctly handles edge cases like avoiding division by zero and ensuring proper device placement.
48-83: Well-documented function for graph index mappingThe
get_graph_indexfunction provides comprehensive documentation of parameters and return values, making it clear how the function constructs indices for edge and angle graphs from neighbor lists.deepmd/pt/model/descriptor/repflows.py (4)
139-150: Well-documented dynamic neighbor selection featureThe documentation clearly explains the purpose and benefits of the dynamic neighbor selection feature, including how it allows for larger selection numbers and removes the need for fixed padding.
201-244: Good parameter validation and dependency checkingThe code properly validates the new parameters: checking that
smooth_edge_updatemust be True whenuse_dynamic_selis True, and ensuringsel_reduce_factoris positive.
297-299: Consistent parameter propagation to layersThe new parameters are correctly propagated to each RepFlowLayer instance, ensuring consistent configuration throughout the model.
599-612: Effective conditional calculation of rotation matrix embeddingThe code properly switches between static and dynamic calculation methods for the rotation matrix embedding, applying the scaling factor for normalization in the dynamic case.
deepmd/pd/model/descriptor/dpa3.py (1)
153-155: Parameter propagation correctly implementedThe new parameters are properly passed from the
repflow_argsto theDescrptBlockRepflowsconstructor.deepmd/pd/model/descriptor/repflows.py (3)
162-164: Parameters added with appropriate defaultsThe new parameters for dynamic neighbor selection are added with the same default values as in the PyTorch implementation.
196-198: Properly disabling unsupported feature with clear error messageThe code explicitly asserts that dynamic selection is not supported in this implementation, preventing users from enabling a feature that isn't fully implemented.
251-253: Parameter propagation to layer instancesThe new parameters are correctly passed to each
RepFlowLayerinstance, preparing for future support of the feature.deepmd/pd/model/descriptor/repflow_layer.py (5)
55-56: New parameters for dynamic neighbor selection look good.These parameters enable the dynamic neighbor selection feature, allowing for variable-sized neighbor lists without fixed padding.
103-106: Well-implemented instance variables for dynamic selection.The scaling factors
dynamic_e_selanddynamic_a_selare properly derived from the original selection sizes, allowing for proper normalization when using dynamic selection.
821-821: Version increment is appropriate for the API change.Properly incrementing the version number from 1 to 2 ensures backward compatibility can be maintained during deserialization.
845-846: New parameters correctly added to serialization.The dynamic selection parameters are properly included in the serialized data structure.
889-889: Correctly updated version compatibility check.The
check_version_compatibilityfunction now correctly handles both version 1 and 2 models during deserialization.deepmd/dpmodel/descriptor/repflows.py (13)
24-25: Added utility imports for dynamic selection functionality.These imported functions support the new dynamic selection feature, providing graph indexing and aggregation capabilities.
128-139: Well-documented new parameters for dynamic selection.The docstring clearly explains the purpose and behavior of the dynamic neighbor selection feature, including the recommendation to set larger selection numbers when this feature is enabled.
221-230: Good validation for dynamic selection parameters.The code properly validates that:
smooth_edge_updatemust be True whenuse_dynamic_selis Truesel_reduce_factormust be greater than 0This prevents invalid configurations that could lead to runtime errors.
483-512: More explicit handling of neighbor list index padding.The code now explicitly sets padding positions to index 0 and adds clear comments, improving code readability and maintainability.
514-536: Clean implementation of dynamic selection logic.The conditional branch for dynamic selection properly:
- Gets graph indices
- Flattens tensors based on masks
- Creates angle masks by combining atomic masks
- Uses zero indices for the static case
This effectively handles both dynamic and static selection modes.
565-578: Rotation matrix calculation handles both modes correctly.The code elegantly branches between static and dynamic rotation matrix calculations, using the appropriate scale factor for normalization in the dynamic case.
721-775: Proper symmetrization operator for dynamic selection mode.The implementation correctly adapts the symmetrization operation to work with flattened dynamic data instead of fixed-size tensors.
1071-1128: Well-implemented angle update for dynamic selection.The dynamic angle update method properly handles flattened tensors and uses index arrays for data lookup and reduction.
1173-1219: Effective edge update for dynamic selection.The dynamic edge update method implements the equivalent operations to the static version but adapts them for flattened data structures.
1435-1450: Elegant normalization for dynamic selection.The normalization in the dynamic case properly uses the
dynamic_e_selvalue instead of the rawnnei, ensuring consistent scaling behavior between static and dynamic modes.
1640-1644: Good error handling for incompatible configurations.The code properly blocks invalid configurations (using dynamic selection without smooth edge updates) with a clear error message explaining the restriction.
1781-1782: New parameters correctly added to serialization structure.The dynamic selection parameters are properly incorporated into the serialization format.
1825-1825: Correctly updated version compatibility check.Version checking now properly handles both version 1 and 2 models during deserialization.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
deepmd/dpmodel/descriptor/repflows.py (1)
666-718:⚠️ Potential issueMissing check for empty edge sets in _cal_hg_dynamic
This function correctly calculates the transposed rotation matrix for dynamic selection, but it doesn't handle the case when
n_edge==0, which can happen with extreme exclusions and cause crashes in the underlying implementation.Add the following early return at the beginning of the function, after line 705:
def _cal_hg_dynamic( flat_edge_ebd: np.ndarray, flat_h2: np.ndarray, flat_sw: np.ndarray, owner: np.ndarray, num_owner: int, nb: int, nloc: int, scale_factor: float, ) -> np.ndarray: xp = array_api_compat.array_namespace(flat_edge_ebd, flat_h2, flat_sw, owner) n_edge, e_dim = flat_edge_ebd.shape + # Handle empty edge set case + if n_edge == 0: + return xp.zeros((nb, nloc, 3, e_dim)) # n_edge x e_dim flat_edge_ebd = flat_edge_ebd * xp.expand_dims(flat_sw, axis=-1)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/dpmodel/descriptor/repflows.py(29 hunks)deepmd/dpmodel/utils/network.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/dpmodel/utils/network.py
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/dpmodel/descriptor/repflows.py
1296-1296: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- 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: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (true)
🔇 Additional comments (9)
deepmd/dpmodel/descriptor/repflows.py (9)
128-139: Implementation for dynamic neighbor selection looks goodThe dynamic selection feature allows for more accurate handling of neighbors by using variable numbers without padding. This is more efficient and accurate than fixed padding.
221-230: Good validation for the feature requirementsThe validation ensures that dynamic selection is used with smooth edge updates and properly validates the reduction factor parameter.
514-535: The dynamic selection branch correctly handles the graph index and tensor flatteningThis implementation properly handles both the standard and dynamic selection approaches, flattening the tensors when using dynamic selection and using masks appropriately.
847-849: Normalization factors look goodThe calculation of
dynamic_e_selanddynamic_a_selfrom the selection counts divided by the reduction factor is appropriate for normalization in the dynamic selection case.
1071-1129: Optimized angle update method for dynamic selection is well structuredThe dynamic version of the angle update method properly handles flattened tensors and index arrays, maintaining a similar logic to the original version but adapted for dynamic selection.
1173-1220: Optimized edge update method for dynamic selection is well implementedThe dynamic version of the edge update method correctly uses indexing operations to achieve the same result as the original approach but with dynamic selection.
1435-1450: Good scaling adjustment for dynamic selectionThe implementation properly adjusts the scaling factor when using dynamic selection by dividing by
dynamic_e_selinstead of the fixednneivalue, ensuring consistent normalization.
1781-1782: Version increment and serialization handled properlyThe serialization of the new parameters and the version increment to 2 ensures backward compatibility and proper handling of the new feature.
1825-1825: Good version compatibility checkThis ensures backward compatibility when deserializing models, allowing older models to be loaded with the new code.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/dpmodel/utils/network.py (1)
962-1027: Well-structured implementation for backend-agnostic aggregationThe
aggregatefunction provides a robust solution for aggregating rows in data tensors according to owner indices. The implementation effectively handles different backend arrays througharray_api_compatand provides fallback implementations for environments lacking native support.Consider using a ternary operator for better readability:
- if hasattr(xp, "bincount"): - bin_count = xp.bincount(owners) - else: - # for array_api_strict - bin_count = bincount(owners) + # for array_api_strict if native bincount is not available + bin_count = xp.bincount(owners) if hasattr(xp, "bincount") else bincount(owners)🧰 Tools
🪛 Ruff (0.11.9)
1005-1009: Use ternary operator
bin_count = xp.bincount(owners) if hasattr(xp, "bincount") else bincount(owners)instead ofif-else-block(SIM108)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/dpmodel/utils/network.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.11.9)
deepmd/dpmodel/utils/network.py
1005-1009: Use ternary operator bin_count = xp.bincount(owners) if hasattr(xp, "bincount") else bincount(owners) instead of if-else-block
(SIM108)
⏰ Context from checks skipped due to timeout of 90000ms (23)
- 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: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- 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 Python (1, 3.9)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (1)
deepmd/dpmodel/utils/network.py (1)
1029-1123:Details
✅ Verification successful
Thorough implementation for graph indexing with proper type handling
The
get_graph_indexfunction correctly computes index mappings for edge and angle graphs from neighbor lists and masks. It properly handles type conversions and efficiently generates the necessary indices for graph-based operations.I noticed you correctly addressed the previous review comment by using
int(xp.sum(xp.astype(nlist_mask, xp.int32)))at line 1073 to convert the array scalar to a Python integer. This change ensures compatibility with JIT and non-NumPy backends.Let's verify that all other array scalar conversions in this function are handled properly:
🏁 Script executed:
#!/bin/bash # Check for other potential array scalar to integer conversions in the function # Look for direct usage of array values in arange or similar functions rg "xp\.arange\([^)]+\)" deepmd/dpmodel/utils/network.py # Look for other array sum operations that might need int conversion rg "xp\.sum\(" deepmd/dpmodel/utils/network.pyLength of output: 382
All scalar conversions in
get_graph_indexare correctI verified that:
nfandnloc(fromnlist.shape) are Python ints when used inxp.arange(nf * nloc, …).nallis a Python int and only scales an array, not used where an int is required.n_edgeis explicitly converted to a Python int viaint(xp.sum(...))before being passed toxp.arange.- No other
xp.sumresults are used in contexts that require Python ints.No changes are needed here.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/dpmodel/utils/network.py (1)
964-1001: Clean implementation of data aggregation functionThe function properly handles aggregation of data by owner indices with good edge case handling.
The division operation in line 999 could be simplified using broadcasting rather than double transposition:
- output = xp.transpose(xp.transpose(output) / bin_count) + output = output / bin_count[:, xp.newaxis]This would make the code more readable and potentially more efficient.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/dpmodel/array_api.py(1 hunks)deepmd/dpmodel/utils/network.py(2 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
deepmd/dpmodel/utils/network.py (2)
deepmd/dpmodel/array_api.py (2)
add_at(97-115)bincount(118-128)deepmd/pt/model/network/utils.py (2)
aggregate(10-45)get_graph_index(49-135)
🔇 Additional comments (3)
deepmd/dpmodel/array_api.py (2)
97-116: Well-implemented array indexing add operation across backendsThis function elegantly handles the differences between NumPy's in-place
add.at, JAX's functional updates, and provides a reasonable fallback implementation for other array types.The code is well-structured with appropriate comments explaining the backend-specific implementations. The fallback implementation using a loop is a reasonable approach, though it may be less efficient for large arrays.
118-128: Effective implementation of bincount with multi-backend supportThe function provides a consistent interface for bincount operations across backends, with appropriate handling of edge cases.
Good job on handling the case where weights are None and ensuring the output has the correct length using
max(minlength, int(xp.max(x)) + 1). The fallback implementation usingadd_atis particularly clean.deepmd/dpmodel/utils/network.py (1)
1004-1098: Robust implementation of graph index mapping with array scalar fixThis function correctly constructs index mappings for complex graph operations, addressing the array scalar issue noted in previous reviews.
Line 1048 properly addresses the previous review comment by explicitly converting the sum to a Python integer:
n_edge = int(xp.sum(xp.astype(nlist_mask, xp.int32)))This ensures compatibility with backends that return array scalars (like JAX) rather than Python integers, preventing type errors when using with
xp.arange.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4754 +/- ##
==========================================
+ Coverage 84.73% 84.78% +0.04%
==========================================
Files 697 698 +1
Lines 67424 67709 +285
Branches 3541 3541
==========================================
+ Hits 57134 57406 +272
- Misses 9160 9170 +10
- Partials 1130 1133 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Summary by CodeRabbit
New Features
use_dynamic_sel) and adjust normalization scaling (sel_reduce_factor).Bug Fixes
Documentation
Tests
Chores