pd: support dpa2 #4418
pd: support dpa2 #4418
Conversation
for more information, see https://pre-commit.ci
…y to coverage newly added code
for more information, see https://pre-commit.ci
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 52
🧹 Outside diff range and nitpick comments (89)
deepmd/pt/loss/__init__.py (1)
25-31: Consider alphabetically ordering the__all__list.While the current ordering works functionally, maintaining an alphabetical order would improve readability and make future modifications easier to manage.
Here's the suggested reordering:
__all__ = [ - "DOSLoss", "DenoiseLoss", + "DOSLoss", "EnergySpinLoss", "EnergyStdLoss", "PropertyLoss", "TaskLoss", "TensorLoss", ]deepmd/jax/env.py (1)
21-21: Consider adding documentation about jax_export usage.Since
jax_exportis now part of the public API, it would be helpful to add a comment explaining its purpose and intended usage within the DeepMD-kit context, especially given its role in the DPA-2 support implementation.__all__ = [ "jax", + # Exposed for model export functionality used in DPA-2 implementation "jax_export", "jnp", "nnx", ]deepmd/pt/model/task/dos.py (1)
Inconsistent version compatibility checks detected across fitting classes
Most fitting classes use version 3, but
deepmd/pt/model/task/polarizability.pyuses version 4. This inconsistency needs to be addressed.
deepmd/pt/model/task/polarizability.py: Uses version 4- All other fitting classes: Use version 3
- deepmd/pd/model/task/ener.py
- deepmd/pd/model/task/invar_fitting.py
- deepmd/pt/model/task/property.py
- deepmd/pt/model/task/invar_fitting.py
- deepmd/pt/model/task/ener.py
- deepmd/pt/model/task/dipole.py
- deepmd/pt/model/task/dos.py
🔗 Analysis chain
Line range hint
104-113: Verify version compatibility across fitting classesThe version compatibility check has been updated from 2 to 3, which is a breaking change.
Let's verify this change is consistent across other fitting classes:
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check version compatibility in other fitting classes # Expected: All fitting classes should use version 3 for consistency # Search for version compatibility checks in fitting classes rg -A 1 "check_version_compatibility" --glob "deepmd/**/fitting*.py" --glob "deepmd/**/task/*.py"Length of output: 2289
deepmd/dpmodel/model/__init__.py (1)
Line range hint
1-37: Consider adding a changelog entrySince this change modifies the public API by exposing
DPModelCommon, it would be helpful to document this change in the changelog to help users understand the implications of this addition.deepmd/dpmodel/fitting/property_fitting.py (1)
81-81: Add documentation for the newdim_case_embdparameter.The new parameter has been added correctly, but its documentation is missing in the class docstring. Please update the docstring to include:
- Purpose and usage of
dim_case_embd- Expected value range
- Default behavior when set to 0
Also applies to: 103-103
deepmd/pt/model/atomic_model/pairtab_atomic_model.py (2)
144-151: Enhance docstring to indicate unsupported operationThe implementation correctly raises NotImplementedError for unsupported case identification. Consider updating the docstring to explicitly mention that this operation is not supported by PairTabAtomicModel.
def set_case_embd(self, case_idx: int): """ - Set the case embedding of this atomic model by the given case_idx, - typically concatenated with the output of the descriptor and fed into the fitting net. + Set the case embedding of this atomic model by the given case_idx. + + Note: This operation is not supported by PairTabAtomicModel as it does not + utilize case identification in its implementation. + + Parameters + ---------- + case_idx : int + The index of the case to be embedded. + + Raises + ------ + NotImplementedError + Always raised as case identification is not supported. """
497-506: Enhance docstring with parameter documentationThe implementation correctly indicates that compression is not supported. Consider enhancing the docstring to document the parameters and explain why compression isn't supported for this model type.
def enable_compression( self, min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Pairtab model does not support compression.""" + """ + Enable model compression (not supported). + + PairTabAtomicModel does not support compression as it already uses a + pre-computed tabulation approach for efficiency. + + Parameters + ---------- + min_nbor_dist : float + Minimum neighbor distance for compression (not used). + table_extrapolate : float, optional + Table extrapolation parameter (not used), by default 5. + table_stride_1 : float, optional + First table stride parameter (not used), by default 0.01. + table_stride_2 : float, optional + Second table stride parameter (not used), by default 0.1. + check_frequency : int, optional + Frequency of compression checks (not used), by default -1. + """deepmd/env.py (1)
17-21: Consider adding docstring to document the public interface.Since these variables are now part of the public interface, it would be helpful to add a module-level docstring explaining their purpose and usage.
Example addition:
+ """ + Global configuration module for DeepMD-kit. + + Exports: + GLOBAL_ENER_FLOAT_PRECISION: NumPy dtype for energy calculations + GLOBAL_NP_FLOAT_PRECISION: NumPy dtype for general calculations + SHARED_LIB_MODULE: Name of the shared library module + global_float_prec: String identifier for precision level ("float" or "double") + """deepmd/pt/model/atomic_model/__init__.py (1)
50-56: LGTM! Consider grouping related model classes together.The changes to
__all__maintain consistency with the imported classes. However, consider grouping related model classes together for better organization. For example:
- Base classes (
BaseAtomicModel)- Core models (
DPAtomicModel)- Property models (
DPPropertyAtomicModel,DPDipoleAtomicModel)- Energy models (
DPEnergyAtomicModel,LinearEnergyAtomicModel,DPZBLLinearEnergyAtomicModel)- Special models (
PairTabAtomicModel)deepmd/pt/model/task/dipole.py (1)
54-55: Enhance documentation for dim_case_embd parameterWhile the parameter is properly documented, it would be helpful to provide more context about what case-specific embedding means and how it affects the dipole fitting network's behavior.
deepmd/pt/cxx_op.py (1)
Line range hint
1-100: Consider unifying error handling across PyTorch and Paddle backends.Since this PR is implementing DPA-2 support in the Paddle backend, it would be beneficial to ensure consistent error handling strategies between PyTorch and Paddle implementations. Consider:
- Implementing similar version and ABI compatibility checks in the Paddle backend
- Using a shared error message formatting strategy across backends
- Creating common utility functions for version compatibility checks
This will help maintain consistency and reduce duplication as the codebase grows with multiple backend support.
🧰 Tools
🪛 Ruff (0.8.0)
50-50: Yoda condition detected
Rewrite as
pt_cxx11_abi_flag != PT_CXX11_ABI_FLAG(SIM300)
deepmd/pd/model/descriptor/se_t_tebd.py (3)
386-387: Remove unused variableenv_matThe
env_matvariable is assigned but never used in the deserialize method.- env_mat = data.pop("env_mat") + data.pop("env_mat") # Remove unused env_mat data🧰 Tools
🪛 Ruff (0.8.0)
387-387: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
🪛 GitHub Check: CodeQL
[notice] 386-386: Unused local variable
Variable env_mat is not used.
823-827: Optimize tensor operations using einsumThe current implementation uses manual tensor operations that could be simplified and potentially optimized using
paddle.einsum.Consider replacing the manual tensor operations with einsum:
- env_ij = ( - rr_i.unsqueeze(2) * rr_j.unsqueeze(1) - ).sum(-1) + env_ij = paddle.einsum("ijm,ikm->ijk", rr_i, rr_j) - res_ij = ( - env_ij.unsqueeze(-1) * gg - ).sum([1, 2]) + res_ij = paddle.einsum("ijk,ijkm->im", env_ij, gg)This change would make the code more readable and potentially more efficient as einsum operations are often optimized at a lower level.
Also applies to: 909-913
709-713: Simplify conditional assignment using ternary operatorThe current if-else block can be simplified using a ternary operator for better readability.
- if callable(merged): - # only get data for once - sampled = merged() - else: - sampled = merged + # only get data for once if callable + sampled = merged() if callable(merged) else merged🧰 Tools
🪛 Ruff (0.8.0)
709-713: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
deepmd/pd/utils/exclude_mask.py (1)
104-104: Remove or clarify the uncertain commentThe comment "may have a better place for this method..." suggests uncertainty about the method placement. If there's a better location for this method, it should be moved there. Otherwise, the comment should be removed to avoid confusion.
deepmd/pd/model/descriptor/repformers.py (3)
441-499: Document or remove commented-out codeThe commented-out code block for spin handling appears to be a work in progress. If this is planned functionality, it should be properly documented with a TODO comment explaining the implementation plan. If not needed, it should be removed.
Consider adding a TODO comment at the start of the block:
# TODO: Implement spin handling functionality # This block will handle: # 1. Padding for spin components # 2. Communication between processes # 3. Virtual particle handling
365-522: Consider breaking down the forward methodThe forward method is quite complex with multiple responsibilities. Consider extracting some logic into helper methods for better maintainability:
- Neighbor list processing
- Environment matrix computation
- Layer processing
- Rotation matrix computation
Example refactor:
def _process_neighbor_list(self, nlist, extended_atype): # Extract neighbor list processing logic def _compute_environment_matrix(self, extended_coord, nlist, atype): # Extract environment matrix computation logic def _process_layers(self, g1, g2, h2, nlist, nlist_mask, sw): # Extract layer processing logic🧰 Tools
🪛 Ruff (0.8.0)
412-412: Local variable
ng1is assigned to but never usedRemove assignment to unused variable
ng1(F841)
435-435: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
🪛 GitHub Check: CodeQL
[notice] 416-416: Unused local variable
Variable ng1 is not used.
548-552: Simplify conditional assignment using ternary operatorThe if-else block can be simplified using a ternary operator for better readability.
- if callable(merged): - # only get data for once - sampled = merged() - else: - sampled = merged + sampled = merged() if callable(merged) else merged🧰 Tools
🪛 Ruff (0.8.0)
548-552: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
deepmd/dpmodel/fitting/dipole_fitting.py (2)
98-98: Add documentation for the new parameterdim_case_embdThe new parameter needs to be documented in the class docstring under the Parameters section. This will help users understand its purpose and usage.
Add this to the class docstring:
Parameters ---------- + dim_case_embd : int, optional + Dimension for case-specific embeddings. Defaults to 0.
Line range hint
1-234: Architecture Review: Consider performance optimization for large-scale computationsThe implementation is well-structured and maintains type safety. However, given this is part of the DPA-2 support in Paddle backend, consider these optimizations:
- The matrix multiplication in the
callmethod could potentially benefit from Paddle's optimized operations- The reshape operations might be optimizable for better memory efficiency
Would you like assistance in exploring these optimizations?
deepmd/pd/model/network/init.py (2)
56-56: Rename ambiguous variablelfor better readabilityThe single-letter variable
lcan be easily confused with the number1or uppercaseI. Consider renaming it to something more descriptive likelower_bound.- l = norm_cdf((a - mean) / std) - u = norm_cdf((b - mean) / std) + lower_bound = norm_cdf((a - mean) / std) + upper_bound = norm_cdf((b - mean) / std) - tensor.uniform_(2 * l - 1, 2 * u - 1) + tensor.uniform_(2 * lower_bound - 1, 2 * upper_bound - 1)🧰 Tools
🪛 Ruff (0.8.0)
56-56: Ambiguous variable name:
l(E741)
328-328: Add explicit stacklevel to warning messagesThe warnings should include an explicit stacklevel to ensure they point to the correct location in the call stack.
- warnings.warn("Initializing zero-element tensors is a no-op") + warnings.warn("Initializing zero-element tensors is a no-op", stacklevel=2)Also applies to: 378-378
🧰 Tools
🪛 Ruff (0.8.0)
328-328: No explicit
stacklevelkeyword argument found(B028)
deepmd/pt/model/task/ener.py (1)
53-53: LGTM! Consider adding parameter documentation.The addition of
dim_case_embdparameter and its propagation to the superclass is well-implemented and maintains backward compatibility. However, the parameter's purpose and usage should be documented in the docstring.Add parameter documentation:
def __init__( self, ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], bias_atom_e: Optional[torch.Tensor] = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, seed: Optional[Union[int, list[int]]] = None, type_map: Optional[list[str]] = None, **kwargs, ) -> None: + """Initialize EnergyFittingNet. + + Args: + dim_case_embd: Dimension of case embedding. Defaults to 0. + [document other parameters...] + """Also applies to: 71-71
deepmd/pt/model/descriptor/dpa2.py (1)
Line range hint
445-519: Add validation for the new type map.While the implementation is thorough, consider adding validation for the new type map to ensure it's well-formed:
- Check for duplicate types
- Validate type names (e.g., non-empty strings)
- Handle empty type maps gracefully
def change_type_map( self, type_map: list[str], model_with_new_type_stat=None ) -> None: assert ( self.type_map is not None ), "'type_map' must be defined when performing type changing!" + # Validate new type map + if not type_map: + raise ValueError("Type map cannot be empty") + if len(set(type_map)) != len(type_map): + raise ValueError("Duplicate types found in type map") + if any(not isinstance(t, str) or not t.strip() for t in type_map): + raise ValueError("Invalid type names in type map") remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map)deepmd/dpmodel/fitting/polarizability_fitting.py (1)
Line range hint
1-324: Consider adding tests for case embedding functionalityWhile the implementation looks good, we should ensure proper test coverage for the new case embedding functionality in the polarizability fitting.
Would you like me to help generate test cases for:
- Serialization/deserialization with case embedding
- Forward pass with non-zero case embedding dimension
- Backward compatibility with zero case embedding dimension
deepmd/infer/model_devi.py (2)
Line range hint
19-42: Consider simplifying overload signaturesThe overload signatures use ellipsis (...) for the
atomicparameter default values, but they could be more explicit.Consider this improvement:
@overload def calc_model_devi_f( fs: np.ndarray, real_f: Optional[np.ndarray] = None, relative: Optional[float] = None, - atomic: Literal[False] = ..., + atomic: Literal[False] = False, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... @overload def calc_model_devi_f( fs: np.ndarray, real_f: Optional[np.ndarray] = None, relative: Optional[float] = None, - atomic: Literal[True] = ..., + atomic: Literal[True] = True, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ...
Based on the gathered evidence, I can now provide a final response:
Model deviation functionality needs to be implemented in Paddle backend
The verification reveals:
- DPA-2 descriptor is properly implemented in Paddle backend (
deepmd/pd/model/descriptor/dpa2.py)- However, model deviation functionality is missing in Paddle backend:
- No model deviation implementation found in Paddle's inference module
- No model deviation tests exist for Paddle backend
- The functionality exists only in TensorFlow backend (
deepmd/tf/infer/model_devi.py)🔗 Analysis chain
Line range hint
1-450: Verify DPA-2 implementation completenessWhile the code quality improvements are good, I don't see direct implementation of DPA-2 support as mentioned in the PR objectives.
Let's verify if DPA-2 related changes are present in other files:
The search results show that DPA-2 implementation exists in both Paddle and PyTorch backends:
deepmd/pd/model/descriptor/dpa2.pydeepmd/pt/model/descriptor/dpa2.pydeepmd/jax/descriptor/dpa2.pyLet me verify if the model deviation functionality is properly implemented in the Paddle backend:
Let me check the model deviation implementation in more detail:
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Search for DPA-2 related implementations echo "Searching for DPA-2 related files..." rg -l "dpa2|DPA2|DPA-2" --type py echo "Searching for Paddle backend changes..." rg -l "paddle" --type py echo "Checking model deviation calculations in Paddle..." rg -A 5 "calc_model_devi" --type pyLength of output: 16006
Script:
#!/bin/bash # Check if model deviation is implemented in Paddle backend echo "Checking model deviation implementation in Paddle backend..." rg -l "model_devi" deepmd/pd/ # Check DPA2 descriptor implementation in Paddle echo -e "\nChecking DPA2 descriptor implementation in Paddle..." rg -l "class.*DPA2" deepmd/pd/model/descriptor/dpa2.py # Compare DPA2 implementations across backends echo -e "\nComparing DPA2 implementations across backends..." for f in deepmd/{pd,pt,jax}/model/descriptor/dpa2.py; do if [ -f "$f" ]; then echo "=== $f ===" head -n 20 "$f" fi doneLength of output: 1855
Script:
#!/bin/bash # Check model deviation implementation in Paddle backend's inference module echo "Checking for model deviation in Paddle inference..." fd -e py infer deepmd/pd/ -x cat {} # Check if there are any model deviation tests for Paddle echo -e "\nChecking for model deviation tests in Paddle..." fd -e py "test.*model.*devi" source/tests/pd/ -x cat {}Length of output: 2272
deepmd/pd/utils/decomp.py (5)
3-7: Consider adding tracking issue for future removalThe comment indicates this file is temporary until Paddle implements these primitives. Consider creating a tracking issue and adding its reference here to ensure this technical debt is properly tracked.
53-64: Standardize reduction mode mappingsThe reduction modes are mapped inconsistently:
- "sum" -> "add"
- "prod" -> "mul"
Consider standardizing these mappings to match either the input mode names or the underlying Paddle operation names.
115-118: Consider caching nonzero indices for repeated operationsThe
nonzerooperation can be expensive for large tensors. If this function is called repeatedly with the same mask, consider caching the computed indices.+# If this function is called frequently with the same mask, consider adding: +# @functools.lru_cache(maxsize=128) +def _compute_mask_coords(mask: paddle.Tensor) -> paddle.Tensor: + return paddle.concat(paddle.nonzero(mask, as_tuple=True), axis=1)
130-134: Consider using paddle operations in dynamic modeInstead of using
np.prod, consider using paddle's native operations for consistency and potential performance benefits on GPU:def numel(x: paddle.Tensor) -> int: if paddle.in_dynamic_mode(): - return np.prod(x.shape) + return paddle.prod(paddle.to_tensor(x.shape)) return paddle.numel(x)
137-137: Fix typo in comment"convinience" should be "convenience"
deepmd/pt/model/descriptor/__init__.py (1)
46-58: LGTM! Nice alphabetical ordering.The
__all__list is now properly alphabetized, which improves readability and maintainability. All exported symbols have corresponding imports.deepmd/common.py (1)
65-69: LGTM! Consider documenting type hints usage.The type exports are correctly added to support static type checking. Consider adding a comment or documentation section explaining how to use these type hints in downstream code.
deepmd/pd/train/wrapper.py (4)
30-34: Fix incomplete docstring.The docstring references a non-existent
configparameter and is missing documentation for all actual parameters.Update the docstring to:
- """Construct a DeePMD model wrapper. - - Args: - - config: The Dict-like configuration with training options. - """ + """Construct a DeePMD model wrapper. + + Args: + model: Single model or dictionary of models for multi-task scenarios + loss: Optional single loss or dictionary of losses + model_params: Optional model parameters dictionary + shared_links: Optional parameter sharing configuration + """
23-29: Add type hints for model_params and shared_links.Missing type hints could lead to runtime errors if incorrect types are passed.
def __init__( self, model: paddle.nn.Layer | dict, loss: paddle.nn.Layer | dict = None, - model_params=None, - shared_links=None, + model_params: dict | None = None, + shared_links: dict | None = None, ) -> None:
103-107: Improve error handling for has_spin attribute.The
getattrcall could silently ignore missing has_spin implementations.- has_spin = getattr(self.model[task_key], "has_spin", False) + has_spin = getattr(self.model[task_key], "has_spin", None) + if has_spin is None: + log.warning(f"Model for task {task_key} does not implement has_spin") + has_spin = False if callable(has_spin): has_spin = has_spin()
142-145: Add error handling for extra state.The method should validate the extra state dictionary structure.
def set_extra_state(self, extra_state: dict): + if not isinstance(extra_state, dict): + raise TypeError("extra_state must be a dictionary") + if "model_params" not in extra_state or "train_infos" not in extra_state: + raise KeyError("extra_state must contain 'model_params' and 'train_infos' keys") self.model_params = extra_state["model_params"] self.train_infos = extra_state["train_infos"] return Nonedeepmd/pd/model/network/network.py (2)
205-229: Add input validation to forward method.The forward method lacks input validation for the device parameter, which could lead to runtime errors.
Consider adding validation:
def forward(self, device: str): + if not isinstance(device, (str, paddle.base.libpaddle.Place)): + raise TypeError(f"Expected device to be str or Place, got {type(device)}") if not self.use_econf_tebd: embed = self.embedding_net( paddle.eye(self.ntypes, dtype=self.prec).to(device=device) )
296-321: Enhance error handling in deserialize method.The deserialize method could benefit from more robust error handling and validation.
Consider adding more comprehensive error handling:
@classmethod def deserialize(cls, data: dict): + if not isinstance(data, dict): + raise TypeError(f"Expected dict, got {type(data)}") + if not data: + raise ValueError("Empty data dictionary provided") + data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 1) data_cls = data.pop("@class") assert data_cls == "TypeEmbedNet", f"Invalid class {data_cls}" + required_keys = {"ntypes", "neuron", "activation_function"} + missing_keys = required_keys - set(data.keys()) + if missing_keys: + raise ValueError(f"Missing required keys: {missing_keys}") + embedding_net = EmbeddingNet.deserialize(data.pop("embedding")) # compat with version 1 if "use_tebd_bias" not in data: data["use_tebd_bias"] = True type_embedding_net = cls(**data) type_embedding_net.embedding_net = embedding_net return type_embedding_netdeepmd/pd/model/task/invar_fitting.py (1)
176-176: Fix typo in docstringThere's a typo in the word "alculate" in the docstring.
- """Based on embedding net output, alculate total energy. + """Based on embedding net output, calculate total energy.deepmd/pd/model/model/__init__.py (2)
75-95: Simplify nested loops and improve error messages.The nested loops and error handling could be more elegant and informative.
Consider this improvement:
def _convert_preset_out_bias_to_array(preset_out_bias, type_map): if preset_out_bias is not None: for kk in preset_out_bias: if len(preset_out_bias[kk]) != len(type_map): raise ValueError( - "length of the preset_out_bias should be the same as the type_map" + f"Length mismatch for '{kk}': expected {len(type_map)}, got {len(preset_out_bias[kk])}" ) - for jj in range(len(preset_out_bias[kk])): - if preset_out_bias[kk][jj] is not None: - if isinstance(preset_out_bias[kk][jj], list): - bb = preset_out_bias[kk][jj] - elif _can_be_converted_to_float(preset_out_bias[kk][jj]): - bb = [float(preset_out_bias[kk][jj])] - else: - raise ValueError( - f"unsupported type/value of the {jj}th element of " - f"preset_out_bias['{kk}'] " - f"{type(preset_out_bias[kk][jj])}" - ) - preset_out_bias[kk][jj] = np.array(bb) + preset_out_bias[kk] = [ + np.array(v if isinstance(v, list) else [float(v)]) + if v is not None and (_can_be_converted_to_float(v) or isinstance(v, list)) + else v + for v in preset_out_bias[kk] + ] return preset_out_bias
129-134: Add type hints and documentation.The function would benefit from type hints and a proper docstring.
Consider this improvement:
-def get_model(model_params): +def get_model(model_params: dict) -> BaseModel: + """Factory method to create model instances. + + Args: + model_params (dict): Model configuration parameters + + Returns: + BaseModel: Created model instance based on the specified type + + Raises: + ValueError: If model_params is invalid + """ + if not isinstance(model_params, dict): + raise ValueError("model_params must be a dictionary") + model_type = model_params.get("type", "standard") if model_type == "standard": return get_standard_model(model_params) else: return BaseModel.get_class_by_type(model_type).get_model(model_params)deepmd/pd/model/descriptor/se_atten.py (4)
375-379: Simplify Conditional Assignment Using Ternary OperatorThe conditional assignment to
sampledcan be simplified using a ternary operator for conciseness.Apply this change:
- if callable(merged): - sampled = merged() - else: - sampled = merged + sampled = merged() if callable(merged) else merged🧰 Tools
🪛 Ruff (0.8.0)
375-379: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
477-477: Remove Assignment to Unused VariablenallThe variable
nallis assigned but never used, which may cause confusion. Consider removing it if it's not needed.Apply this change:
- nall = extended_coord.reshape([nb, -1, 3]).shape[1]🧰 Tools
🪛 Ruff (0.8.0)
477-477: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
467-470: Clarify Return Types in the DocstringIn the
forwardmethod's docstring (lines 467-470), the description of the return values could be more precise. Consider specifying the types and shapes of the returned tensors for better clarity.Update the docstring as follows:
Returns ------- result : paddle.Tensor - The descriptor. shape: nf x nloc x (ng x axis_neuron) + The descriptor tensor with shape: [nf, nloc, ng * axis_neuron] g2 : paddle.Tensor - The rotationally invariant pair-partical representation. - shape: nf x nloc x nnei x ng + Rotationally invariant pair-particle representation with shape: [nf, nloc, nnei, ng] h2 : paddle.Tensor - The rotationally equivariant pair-partical representation. - shape: nf x nloc x nnei x 3 + Rotationally equivariant pair-particle representation with shape: [nf, nloc, nnei, 3] gr : paddle.Tensor - The rotationally equivariant and permutationally invariant single particle - representation. shape: nf x nloc x ng x 3 + Rotationally equivariant and permutationally invariant single-particle representation with shape: [nf, nloc, ng, 3] sw : paddle.Tensor - The smooth switch function. shape: nf x nloc x nnei + Smooth switching function with shape: [nf, nloc, nnei]
587-588: Implement Compression Feature or Remove Unimplemented CodeThe method raises a
NotImplementedErrorat lines 587-588, indicating that compression is not implemented yet. If compression is a planned feature, consider implementing it or providing a more descriptive error message. If not, consider removing this code block to avoid runtime errors.deepmd/dpmodel/atomic_model/pairtab_atomic_model.py (1)
438-447: Consider improving the compression method implementation.For consistency with
set_case_embd, consider:
- Raising NotImplementedError instead of silently passing
- Expanding the docstring to explain why compression isn't supported for this model
def enable_compression( self, min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: - """Pairtab model does not support compression.""" - pass + """ + The PairTab model does not support compression as it already uses + a pre-computed tabulation for efficiency. + """ + raise NotImplementedError( + "Compression not supported for PairTabAtomicModel as it uses pre-computed tabulation!" + )deepmd/pd/model/descriptor/dpa2.py (5)
499-499: Rename unused loop variableiito_for clarityIn the loop at line 499, the variable
iiis not used within the loop body. Renaming it to_indicates that the variable is intentionally unused.Apply this diff to rename the unused variable:
-for ii, descrpt in enumerate(descrpt_list): +for _, descrpt in enumerate(descrpt_list): descrpt.compute_input_stats(merged, path)🧰 Tools
🪛 Ruff (0.8.0)
499-499: Loop control variable
iinot used within loop bodyRename unused
iito_ii(B007)
653-653: Remove unused variableenv_matinrepinitdeserializationThe variable
env_matis assigned but not used afterward in the deserialization ofrepinit. Removing this assignment will clean up the code.Apply this diff to remove the unused assignment:
- env_mat = repinit_variable.pop("env_mat")
668-668: Remove unused variableenv_matinrepinit_three_bodydeserializationThe variable
env_matis assigned but not used afterward in the deserialization ofrepinit_three_body. Removing this assignment will tidy up the code.Apply this diff to remove the unused assignment:
- env_mat = repinit_three_body_variable.pop("env_mat")
684-684: Remove unused variableenv_matinrepformersdeserializationThe variable
env_matis assigned but not used afterward in the deserialization ofrepformers. Removing this assignment will improve code clarity.Apply this diff to remove the unused assignment:
- env_mat = repformers_variable.pop("env_mat")🧰 Tools
🪛 Ruff (0.8.0)
684-684: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
873-897: Implement theenable_compressionmethod or remove the placeholderThe
enable_compressionmethod currently raisesNotImplementedError. If compression is a planned feature, consider implementing this method. Otherwise, remove the placeholder to avoid confusion and potential runtime errors.Would you like assistance in implementing the
enable_compressionmethod or should we create a GitHub issue to track this task?deepmd/pd/model/task/fitting.py (4)
419-448: Improve error messages and add performance optimization in _forward_commonThe error messages could be more descriptive, and there's an opportunity for performance optimization in tensor operations.
Consider these improvements:
def _forward_common(self, descriptor: paddle.Tensor, atype: paddle.Tensor, ...): + # Validate input shapes early + if descriptor.ndim != 3: + raise ValueError( + f"descriptor must be 3D tensor, got shape {descriptor.shape}" + ) + if atype.ndim != 2: + raise ValueError(f"atype must be 2D tensor, got shape {atype.shape}") + # cast the input to internal precsion - xx = descriptor.to(self.prec) - fparam = fparam.to(self.prec) if fparam is not None else None - aparam = aparam.to(self.prec) if aparam is not None else None + # Batch convert tensors to improve performance + tensors_to_convert = [t for t in [descriptor, fparam, aparam] if t is not None] + converted_tensors = [t.to(self.prec) for t in tensors_to_convert] + xx, fparam, aparam = ( + converted_tensors[0], + converted_tensors[1] if fparam is not None else None, + converted_tensors[2] if aparam is not None else None, + )
328-331: Optimize dictionary key lookupThe key lookup in the deserialize method can be optimized by removing the
.keys()call.- for kk in variables.keys(): + for kk in variables: obj[kk] = to_paddle_tensor(variables[kk])🧰 Tools
🪛 Ruff (0.8.0)
328-328: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
281-320: Document version compatibility in serialize methodThe serialize method includes version information but lacks documentation about version compatibility and breaking changes.
Add docstring to explain version compatibility:
def serialize(self) -> dict: - """Serialize the fitting to dict.""" + """Serialize the fitting to dict. + + Returns + ------- + dict + A dictionary containing the serialized data. + Version compatibility: + - Version 1: Initial version + - Version 2: Added support for mixed types + - Version 3: Added support for case embeddings + """
1-548: Overall well-structured implementation with room for improvementsThe code implements a complex fitting network with good architecture and separation of concerns. While the core functionality is solid, consider these general improvements:
- Add comprehensive docstrings for all public methods
- Implement logging for better debugging
- Add unit tests for edge cases
- Consider adding type hints for internal methods
Would you like assistance in implementing any of these improvements?
🧰 Tools
🪛 Ruff (0.8.0)
132-132: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
143-143: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
251-251: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
328-328: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
deepmd/pd/model/task/ener.py (1)
34-87: Implementation aligns with DPA-2 support objectivesThe
EnergyFittingNetimplementation follows the established patterns of the fitting framework while adding support for case embeddings throughdim_case_embd. This aligns well with the PR's objective of implementing DPA-2 support in the Paddle backend.A few architectural considerations:
- The class properly extends
InvarFittingand maintains the fitting type registration pattern- Case embedding support is consistently implemented through constructor parameters
- Serialization methods maintain proper versioning for future compatibility
🧰 Tools
🪛 Ruff (0.8.0)
40-40: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/pd/model/descriptor/repformer_layer.py (5)
41-82: Consider adding parameter validation for _scale and _dimThe function handles initialization modes well, but should validate that
_scaleis non-negative and_dimis positive to prevent potential issues.def get_residual( _dim: int, _scale: float, _mode: str = "norm", trainable: bool = True, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, ) -> paddle.Tensor: + if _dim <= 0: + raise ValueError("_dim must be positive") + if _scale < 0: + raise ValueError("_scale must be non-negative") random_generator = get_generator(seed)
85-116: Consider using paddle.gather for better performanceThe current implementation uses
take_along_axiswhich is less efficient thangatherfor this use case.- gg1 = paddle.take_along_axis(g1_ext, axis=1, indices=index) + gg1 = paddle.gather(g1_ext, index, axis=1)
967-967: Remove unused variable assignmentThe variable
ng2is assigned but never used in this context.- nb, nloc, nnei, _ = g2.shape - ng2 = g2.shape[-1] + nb, nloc, nnei, _ = g2.shape🧰 Tools
🪛 Ruff (0.8.0)
967-967: Local variable
ng2is assigned to but never usedRemove assignment to unused variable
ng2(F841)
🪛 GitHub Check: CodeQL
[notice] 967-967: Unused local variable
Variable ng2 is not used.
1157-1160: Simplify conditional assignment using ternary operatorThe if-else block can be simplified using a ternary operator.
- if cal_gg1: - gg1 = _make_nei_g1(g1_ext, nlist) - else: - gg1 = None + gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None🧰 Tools
🪛 Ruff (0.8.0)
1157-1160: Use ternary operator
gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else Noneinstead ofif-else-blockReplace
if-else-block withgg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None(SIM108)
1306-1408: Add validation for required fields in serializationThe serialization method should validate that all required fields are present before creating the dictionary.
def serialize(self) -> dict: + required_fields = ['rcut', 'rcut_smth', 'sel', 'ntypes', 'g1_dim', 'g2_dim'] + for field in required_fields: + if not hasattr(self, field): + raise ValueError(f"Required field {field} is missing") data = { "@class": "RepformerLayer", "@version": 2, # ... rest of the serialization }deepmd/pd/utils/region.py (1)
89-91: Remove or address commented-out code referencing undefined_to_face_distanceThe commented-out code references an undefined function
_to_face_distance, which may cause confusion or indicate incomplete code.Consider removing the commented code or implementing
_to_face_distanceif needed. Apply this diff:-# b_to_face_distance = paddle.vmap( -# _to_face_distance, in_dims=(0), out_dims=(0))deepmd/pd/utils/dataloader.py (4)
285-290: Simplify nestedifstatements to improve readabilityThe nested
ifstatements from lines 285 to 290 can be combined into a singleifstatement using logicalandoperators. This will enhance code readability by reducing indentation levels.Apply this diff to combine the
ifstatements:- if self._queue.qsize() < min(2, max(1, self._queue.maxsize // 2)): - if time.time() - self.start_time > 5 * 60: - if ( - self.warning_time is None - or time.time() - self.warning_time > 15 * 60 - ): + if ( + self._queue.qsize() < min(2, max(1, self._queue.maxsize // 2)) + and time.time() - self.start_time > 5 * 60 + and ( + self.warning_time is None + or time.time() - self.warning_time > 15 * 60 + ) + ):🧰 Tools
🪛 Ruff (0.8.0)
285-290: Use a single
ifstatement instead of nestedifstatements(SIM102)
286-290: Use a single
ifstatement instead of nestedifstatementsCombine
ifstatements usingand(SIM102)
310-310: Iterate directly over the dictionary for efficiencyInstead of using
for key in example.keys():, you can iterate directly over the dictionary withfor key in example:. This is more Pythonic and slightly more efficient.Apply this diff to modify the iteration:
- for key in example.keys(): + for key in example:🧰 Tools
🪛 Ruff (0.8.0)
310-310: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
318-319: Add a comment to explain skipping the 'type' keyIn the
collate_batchfunction, the code skips processing the'type'key. Adding a comment to explain the reasoning will improve code clarity and maintainability.Apply this diff to add a comment:
elif key == "type": + # Skip 'type' key as it does not require collation continue
98-108: Renameprocessesvariable tothreadsfor claritySince
multiprocessing.dummy.Pooluses threads rather than processes, renaming theprocessesvariable tothreadswill prevent confusion and enhance code readability.Apply this diff to rename the variable:
- MAX_PROCESSES_NUM = 4 - processes = min( + MAX_THREADS_NUM = 4 + threads = min( os.cpu_count() // ( dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1 ), - MAX_PROCESSES_NUM, + MAX_THREADS_NUM, ) - with Pool(processes) as pool: + with Pool(threads) as pool:deepmd/pt/model/task/fitting.py (1)
215-215: Consider removing the commented-out code for clarity.The commented line initializing
case_embdis not needed and can be removed to keep the code clean.deepmd/pd/utils/env.py (2)
72-76: Remove unused variables and commented-out code.The variables
inter_nthreadsandintra_nthreadsare assigned but not used. Additionally, there is commented-out code that appears to be inactive. Removing unused code improves maintainability and readability.Apply this diff to clean up the code:
set_default_nthreads() -inter_nthreads, intra_nthreads = get_default_nthreads() -# if inter_nthreads > 0: # the behavior of 0 is not documented -# os.environ['OMP_NUM_THREADS'] = str(inter_nthreads) -# if intra_nthreads > 0: -# os.environ['CPU_NUM'] = str(intra_nthreads)If you intend to use these variables or code in the future, consider adding a comment explaining their purpose or implement them accordingly.
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 72-72: Unused global variable
The global variable 'inter_nthreads' is not used.
[notice] 72-72: Unused global variable
The global variable 'intra_nthreads' is not used.
[notice] 75-76: Commented-out code
This comment appears to contain commented-out code.
15-16: Remove unused importget_default_nthreads.After removing the unused variables and code, the function
get_default_nthreadsis no longer used in this module. Consider removing it from the imports to clean up the code.Apply this diff to remove the unused import:
from deepmd.env import ( GLOBAL_ENER_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION, - get_default_nthreads, set_default_nthreads, )deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
49-54: Consider adding type hints fordescriptorandfittingparametersAdding type hints to the
descriptorandfittingparameters in the__init__method enhances code readability and helps with static type checking. It also makes the expected types clear to other developers.def __init__( - self, - descriptor, - fitting, - type_map: list[str], - **kwargs, +) -> None: + self, + descriptor: BaseDescriptor, + fitting: BaseFitting, + type_map: list[str], + **kwargs, ) -> None:deepmd/pd/model/descriptor/env_mat.py (1)
33-34: Improve Variable Naming for ClarityVariables like
t0,t1, and_env_mat_se_aare not descriptive, which can affect code readability and maintainability. Renaming them to more meaningful names would make the code clearer.Example of improved variable names:
# Calculations in _make_env_mat - t0 = 1 / (length + protection) - t1 = diff / (length + protection) ** 2 + inverse_length = 1 / (length + protection) + normalized_diff = diff / (length + protection) ** 2 # ... if radial_only: - env_mat = t0 * weight + env_mat = inverse_length * weight else: - env_mat = paddle.concat([t0.astype(t1.dtype), t1], axis=-1) * weight + env_mat = paddle.concat([inverse_length.astype(normalized_diff.dtype), normalized_diff], axis=-1) * weight # ... # Normalization in prod_env_mat - env_mat_se_a = (_env_mat_se_a - t_avg) / t_std + normalized_env_mat = (env_mat - t_avg) / t_stdAlso applies to: 81-81
deepmd/pd/entrypoints/main.py (2)
116-118: Remove unused variablefThe variable
fis assigned but never used. Since you're creating the file without needing to interact with it, you can omit assigning it tof.Apply this diff to remove the unused variable:
if stat_file_path_single.endswith((".h5", ".hdf5")): - with h5py.File(stat_file_path_single, "w") as f: + with h5py.File(stat_file_path_single, "w"): pass🧰 Tools
🪛 Ruff (0.8.0)
118-118: Local variable
fis assigned to but never usedRemove assignment to unused variable
f(F841)
500-503: Simplify conditional assignment ofFLAGSYou can use a ternary operator to simplify the assignment of
FLAGS.Apply this diff to simplify the code:
- if not isinstance(args, argparse.Namespace): - FLAGS = parse_args(args=args) - else: - FLAGS = args + FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args🧰 Tools
🪛 Ruff (0.8.0)
500-503: Use ternary operator
FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else argsinstead ofif-else-blockReplace
if-else-block withFLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args(SIM108)
deepmd/pd/model/descriptor/se_a.py (3)
47-52: Simplify import handling usingcontextlib.suppressReplace the
try-except-passblock withcontextlib.suppressfor cleaner code.Apply this diff:
+ from contextlib import suppress # ... - try: - from typing import ( - Final, - ) - except ImportError: - pass + with suppress(ImportError): + from typing import Final🧰 Tools
🪛 Ruff (0.8.0)
47-52: Use
contextlib.suppress(ImportError)instead oftry-except-passReplace with
contextlib.suppress(ImportError)(SIM105)
🪛 GitHub Check: CodeQL
[notice] 50-50: Empty except
'except' clause does nothing but pass and there is no explanatory comment.
361-361: Remove unused variableenv_matThe variable
env_matis assigned but never used. You can remove this assignment to clean up the code.Apply this diff:
embeddings = data.pop("embeddings") - env_mat = data.pop("env_mat") obj = cls(**data)🧰 Tools
🪛 Ruff (0.8.0)
361-361: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
645-650: Remove redundant assignment totiThe variable
tiis assigned a value but then immediately reassigned before being used. Remove the initial assignment totito simplify the code.Apply this diff:
if self.type_one_side: ii = embedding_idx - ti = -1 # paddle.jit is not happy with slice(None) # ti_mask = paddle.ones(nfnl, dtype=paddle.bool, device=dmatrix.place) # applying a mask seems to cause performance degradation ti_mask = None else: # ti: center atom type, ii: neighbor type... ii = embedding_idx // self.ntypes ti = embedding_idx % self.ntypesdeepmd/pd/train/training.py (4)
290-290: Simplify dictionarygetcalls by omitting redundantNoneWhen using
dict.get(key, default), you can omit thedefaultif it'sNone, as it's the default value.Apply this diff:
290,392 - if config.get("learning_rate_dict", None) is not None: + if config.get("learning_rate_dict") is not None:Also applies to: 392-392
🧰 Tools
🪛 Ruff (0.8.0)
290-290: Use
config.get("learning_rate_dict")instead ofconfig.get("learning_rate_dict", None)Replace
config.get("learning_rate_dict", None)withconfig.get("learning_rate_dict")(SIM910)
493-493: Simplify key existence check in dictionaryYou can simplify the key existence check by removing
.keys().Apply this diff:
- if key in _random_state_dict.keys(): + if key in _random_state_dict:🧰 Tools
🪛 Ruff (0.8.0)
493-493: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
863-867: Remove unused variablemoduleThe variable
moduleis assigned but never used. You can remove this assignment to clean up the code.Apply this diff:
- module = ( - self.wrapper.module - if dist.is_available() and dist.is_initialized() - else self.wrapper - ) self.save_model(self.latest_model, lr=cur_lr, step=_step_id)🧰 Tools
🪛 Ruff (0.8.0)
863-863: Local variable
moduleis assigned to but never usedRemove assignment to unused variable
module(F841)
1048-1048: Simplify iteration over dictionaryWhen iterating over a dictionary, you can iterate directly over it without calling
.keys().Apply this diff:
- for key in batch_data.keys(): + for key in batch_data:🧰 Tools
🪛 Ruff (0.8.0)
1048-1048: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
deepmd/pd/utils/utils.py (3)
35-57: Consider using an enum for activation function types and optimize string comparisons.The activation function implementation could be improved in terms of maintainability and performance.
Consider these improvements:
+from enum import Enum, auto + +class ActivationType(Enum): + RELU = auto() + GELU = auto() + TANH = auto() + RELU6 = auto() + SOFTPLUS = auto() + SIGMOID = auto() + LINEAR = auto() + class ActivationFn(paddle.nn.Layer): def __init__(self, activation: str | None): super().__init__() - self.activation: str = activation if activation is not None else "linear" + self.activation: str = (activation or "linear").lower() def forward(self, x: paddle.Tensor) -> paddle.Tensor: """Returns the tensor after applying activation function corresponding to `activation`.""" - if self.activation.lower() == "relu": + if self.activation == "relu": return F.relu(x) - elif self.activation.lower() == "gelu" or self.activation.lower() == "gelu_tf": + elif self.activation in ("gelu", "gelu_tf"): return F.gelu(x, approximate=True) - elif self.activation.lower() == "tanh": + elif self.activation == "tanh": return paddle.tanh(x) - elif self.activation.lower() == "relu6": + elif self.activation == "relu6": return F.relu6(x) - elif self.activation.lower() == "softplus": + elif self.activation == "softplus": return F.softplus(x) - elif self.activation.lower() == "sigmoid": + elif self.activation == "sigmoid": return F.sigmoid(x) - elif self.activation.lower() == "linear" or self.activation.lower() == "none": + elif self.activation in ("linear", "none"): return x else: - raise RuntimeError(f"activation function {self.activation} not supported") + supported = ["relu", "gelu", "tanh", "relu6", "softplus", "sigmoid", "linear", "none"] + raise RuntimeError(f"Activation function '{self.activation}' not supported. Supported functions: {supported}")🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 55-55: Statement has no effect
This statement has no effect.
116-127: Simplify device management logic and add device validation.The device management function could be more concise and robust.
Consider this improvement:
def dict_to_device(sample_dict): + valid_devices = {"cpu", "gpu", "xpu"} # Add other valid devices as needed + if not any(DEVICE.startswith(dev) for dev in valid_devices): + raise ValueError(f"Invalid device: {DEVICE}") + for key in sample_dict: - if isinstance(sample_dict[key], list): - sample_dict[key] = [item.to(DEVICE) for item in sample_dict[key]] - if isinstance(sample_dict[key], np.float32): - sample_dict[key] = ( - paddle.ones(1, dtype=paddle.float32).to(device=DEVICE) - * sample_dict[key] - ) - else: - if sample_dict[key] is not None: - sample_dict[key] = sample_dict[key].to(DEVICE) + value = sample_dict[key] + if value is None: + continue + if isinstance(value, list): + sample_dict[key] = [item.to(DEVICE) for item in value] + elif isinstance(value, np.float32): + sample_dict[key] = paddle.ones(1, dtype=paddle.float32).to(device=DEVICE) * value + else: + sample_dict[key] = value.to(DEVICE)
188-198: Enhance error handling in profiling context manager.The profiling context manager could benefit from better error handling and logging.
Consider this improvement:
+import logging + +logger = logging.getLogger(__name__) + @contextmanager def nvprof_context(enable_profiler: bool, name: str): if enable_profiler: - core.nvprof_nvtx_push(name) + try: + core.nvprof_nvtx_push(name) + logger.debug(f"Started profiling section: {name}") + except Exception as e: + logger.error(f"Failed to start profiling section {name}: {str(e)}") + raise try: yield finally: if enable_profiler: - core.nvprof_nvtx_pop() + try: + core.nvprof_nvtx_pop() + logger.debug(f"Ended profiling section: {name}") + except Exception as e: + logger.error(f"Failed to end profiling section {name}: {str(e)}")deepmd/entrypoints/test.py (4)
707-709: Consider using pathlib's stem and suffix methods for file path manipulation.The current implementation directly concatenates strings for file suffixes. While this works, using pathlib's built-in methods would be more idiomatic and safer.
- detail_path.with_suffix(f".dos.out.{ii}"), + detail_path.with_stem(f"{detail_path.stem}.dos.{ii}").with_suffix(".out"),
721-723: Maintain consistency with the previous file naming pattern.The atomic DOS output follows a similar pattern to the regular DOS output. For consistency, consider applying the same refactoring suggestion.
- detail_path.with_suffix(f".ados.out.{ii}"), + detail_path.with_stem(f"{detail_path.stem}.ados.{ii}").with_suffix(".out"),
863-865: Maintain consistent file naming pattern across test functions.The property output follows the same pattern as DOS output. For consistency, consider applying the same refactoring.
- detail_path.with_suffix(f".property.out.{ii}"), + detail_path.with_stem(f"{detail_path.stem}.property.{ii}").with_suffix(".out"),
877-879: Ensure consistent file naming pattern for atomic properties.The atomic property output should follow the same pattern as other atomic outputs.
- detail_path.with_suffix(f".aproperty.out.{ii}"), + detail_path.with_stem(f"{detail_path.stem}.aproperty.{ii}").with_suffix(".out"),
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (82)
.github/workflows/test_cuda.yml(1 hunks).pre-commit-config.yaml(2 hunks)deepmd/__init__.py(1 hunks)deepmd/common.py(2 hunks)deepmd/dpmodel/__init__.py(1 hunks)deepmd/dpmodel/atomic_model/__init__.py(1 hunks)deepmd/dpmodel/atomic_model/dp_atomic_model.py(1 hunks)deepmd/dpmodel/atomic_model/linear_atomic_model.py(2 hunks)deepmd/dpmodel/atomic_model/make_base_atomic_model.py(1 hunks)deepmd/dpmodel/atomic_model/pairtab_atomic_model.py(2 hunks)deepmd/dpmodel/common.py(1 hunks)deepmd/dpmodel/descriptor/__init__.py(1 hunks)deepmd/dpmodel/fitting/__init__.py(1 hunks)deepmd/dpmodel/fitting/dipole_fitting.py(3 hunks)deepmd/dpmodel/fitting/dos_fitting.py(3 hunks)deepmd/dpmodel/fitting/ener_fitting.py(3 hunks)deepmd/dpmodel/fitting/general_fitting.py(9 hunks)deepmd/dpmodel/fitting/invar_fitting.py(3 hunks)deepmd/dpmodel/fitting/polarizability_fitting.py(4 hunks)deepmd/dpmodel/fitting/property_fitting.py(3 hunks)deepmd/dpmodel/model/__init__.py(1 hunks)deepmd/dpmodel/model/make_model.py(2 hunks)deepmd/dpmodel/utils/__init__.py(1 hunks)deepmd/entrypoints/test.py(4 hunks)deepmd/env.py(1 hunks)deepmd/infer/__init__.py(1 hunks)deepmd/infer/deep_eval.py(2 hunks)deepmd/infer/model_devi.py(1 hunks)deepmd/jax/descriptor/__init__.py(1 hunks)deepmd/jax/env.py(1 hunks)deepmd/jax/fitting/__init__.py(1 hunks)deepmd/jax/model/__init__.py(1 hunks)deepmd/pd/__init__.py(1 hunks)deepmd/pd/entrypoints/main.py(1 hunks)deepmd/pd/loss/ener.py(1 hunks)deepmd/pd/model/atomic_model/dp_atomic_model.py(1 hunks)deepmd/pd/model/descriptor/__init__.py(1 hunks)deepmd/pd/model/descriptor/dpa1.py(1 hunks)deepmd/pd/model/descriptor/dpa2.py(1 hunks)deepmd/pd/model/descriptor/env_mat.py(1 hunks)deepmd/pd/model/descriptor/repformer_layer.py(1 hunks)deepmd/pd/model/descriptor/repformers.py(1 hunks)deepmd/pd/model/descriptor/se_a.py(1 hunks)deepmd/pd/model/descriptor/se_atten.py(1 hunks)deepmd/pd/model/descriptor/se_t_tebd.py(1 hunks)deepmd/pd/model/model/__init__.py(1 hunks)deepmd/pd/model/model/ener_model.py(1 hunks)deepmd/pd/model/model/make_model.py(1 hunks)deepmd/pd/model/network/init.py(1 hunks)deepmd/pd/model/network/network.py(1 hunks)deepmd/pd/model/task/__init__.py(1 hunks)deepmd/pd/model/task/ener.py(1 hunks)deepmd/pd/model/task/fitting.py(1 hunks)deepmd/pd/model/task/invar_fitting.py(1 hunks)deepmd/pd/train/training.py(1 hunks)deepmd/pd/train/wrapper.py(1 hunks)deepmd/pd/utils/__init__.py(1 hunks)deepmd/pd/utils/dataloader.py(1 hunks)deepmd/pd/utils/decomp.py(1 hunks)deepmd/pd/utils/env.py(1 hunks)deepmd/pd/utils/exclude_mask.py(1 hunks)deepmd/pd/utils/nlist.py(1 hunks)deepmd/pd/utils/region.py(1 hunks)deepmd/pd/utils/utils.py(1 hunks)deepmd/pt/cxx_op.py(1 hunks)deepmd/pt/loss/__init__.py(1 hunks)deepmd/pt/model/atomic_model/__init__.py(1 hunks)deepmd/pt/model/atomic_model/dp_atomic_model.py(1 hunks)deepmd/pt/model/atomic_model/linear_atomic_model.py(2 hunks)deepmd/pt/model/atomic_model/pairtab_atomic_model.py(2 hunks)deepmd/pt/model/descriptor/__init__.py(1 hunks)deepmd/pt/model/descriptor/dpa2.py(1 hunks)deepmd/pt/model/model/__init__.py(3 hunks)deepmd/pt/model/model/make_model.py(1 hunks)deepmd/pt/model/network/init.py(1 hunks)deepmd/pt/model/task/__init__.py(1 hunks)deepmd/pt/model/task/dipole.py(4 hunks)deepmd/pt/model/task/dos.py(3 hunks)deepmd/pt/model/task/ener.py(3 hunks)deepmd/pt/model/task/fitting.py(11 hunks)deepmd/pt/model/task/invar_fitting.py(4 hunks)deepmd/pt/model/task/polarizability.py(5 hunks)
⛔ Files not processed due to max files limit (15)
- deepmd/pt/model/task/property.py
- deepmd/pt/train/training.py
- deepmd/pt/train/wrapper.py
- deepmd/pt/utils/init.py
- deepmd/pt/utils/dataloader.py
- deepmd/pt/utils/env.py
- deepmd/pt/utils/plugin.py
- deepmd/tf/init.py
- deepmd/tf/common.py
- deepmd/tf/descriptor/init.py
- deepmd/tf/entrypoints/init.py
- deepmd/tf/entrypoints/main.py
- deepmd/tf/env.py
- deepmd/tf/fit/init.py
- deepmd/tf/fit/dipole.py
✅ Files skipped from review due to trivial changes (4)
- deepmd/jax/fitting/init.py
- deepmd/dpmodel/init.py
- deepmd/init.py
- deepmd/dpmodel/common.py
🚧 Files skipped from review as they are similar to previous changes (7)
- deepmd/pd/init.py
- deepmd/pd/model/task/init.py
- deepmd/dpmodel/model/make_model.py
- deepmd/pd/utils/init.py
- .github/workflows/test_cuda.yml
- deepmd/pd/model/descriptor/init.py
- deepmd/pd/loss/ener.py
🧰 Additional context used
🪛 Ruff (0.8.0)
deepmd/pd/model/model/make_model.py
361-361: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
362-362: Use key not in dict instead of key not in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/descriptor/se_t_tebd.py
129-129: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
136-136: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
387-387: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
454-454: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
521-521: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
528-528: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
709-713: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
733-733: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
787-787: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/utils/dataloader.py
54-54: Use capitalized environment variable FLAGS_CUDNN_DETERMINISTIC instead of FLAGS_cudnn_deterministic
Replace FLAGS_cudnn_deterministic with FLAGS_CUDNN_DETERMINISTIC
(SIM112)
86-86: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
285-290: Use a single if statement instead of nested if statements
(SIM102)
286-290: Use a single if statement instead of nested if statements
Combine if statements using and
(SIM102)
310-310: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/task/ener.py
40-40: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/utils/region.py
27-27: Local variable e is assigned to but never used
Remove assignment to unused variable e
(F841)
deepmd/pd/entrypoints/main.py
118-118: Local variable f is assigned to but never used
Remove assignment to unused variable f
(F841)
500-503: Use ternary operator FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args instead of if-else-block
Replace if-else-block with FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args
(SIM108)
deepmd/pd/utils/exclude_mask.py
17-17: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
25-25: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
70-70: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
78-78: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/model/task/invar_fitting.py
93-93: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
104-104: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/utils/nlist.py
386-386: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/model/descriptor/dpa2.py
90-90: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
499-499: Loop control variable ii not used within loop body
Rename unused ii to _ii
(B007)
684-684: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
deepmd/pd/train/training.py
290-290: Use config.get("learning_rate_dict") instead of config.get("learning_rate_dict", None)
Replace config.get("learning_rate_dict", None) with config.get("learning_rate_dict")
(SIM910)
392-392: Use config.get("learning_rate_dict") instead of config.get("learning_rate_dict", None)
Replace config.get("learning_rate_dict", None) with config.get("learning_rate_dict")
(SIM910)
493-493: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
636-636: Use a context manager for opening files
(SIM115)
646-646: Use a context manager for opening files
(SIM115)
683-686: Use ternary operator pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr instead of if-else-block
Replace if-else-block with pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr
(SIM108)
739-739: Loop control variable ii not used within loop body
Rename unused ii to _ii
(B007)
863-863: Local variable module is assigned to but never used
Remove assignment to unused variable module
(F841)
1048-1048: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/task/fitting.py
132-132: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
143-143: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
251-251: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
328-328: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/descriptor/repformer_layer.py
967-967: Local variable ng2 is assigned to but never used
Remove assignment to unused variable ng2
(F841)
1157-1160: Use ternary operator gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None instead of if-else-block
Replace if-else-block with gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None
(SIM108)
1278-1278: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
deepmd/pd/model/descriptor/repformers.py
81-81: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
360-360: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
412-412: Local variable ng1 is assigned to but never used
Remove assignment to unused variable ng1
(F841)
435-435: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
548-552: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
deepmd/pd/model/descriptor/se_atten.py
61-61: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
78-78: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
375-379: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
399-399: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
477-477: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/model/descriptor/dpa1.py
218-218: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
230-230: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
534-534: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
633-633: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/model/descriptor/se_a.py
47-52: Use contextlib.suppress(ImportError) instead of try-except-pass
Replace with contextlib.suppress(ImportError)
(SIM105)
76-76: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
82-82: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
226-226: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
361-361: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
414-414: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
420-420: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
607-611: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
631-631: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
643-643: Loop control variable ll not used within loop body
Rename unused ll to _ll
(B007)
717-717: Loop control variable compress_data_ii not used within loop body
Rename unused compress_data_ii to _compress_data_ii
(B007)
717-717: Loop control variable compress_info_ii not used within loop body
Rename unused compress_info_ii to _compress_info_ii
(B007)
deepmd/pd/model/network/init.py
56-56: Ambiguous variable name: l
(E741)
328-328: No explicit stacklevel keyword argument found
(B028)
378-378: No explicit stacklevel keyword argument found
(B028)
🪛 GitHub Check: CodeQL
deepmd/pd/model/model/make_model.py
[notice] 323-328: Commented-out code
This comment appears to contain commented-out code.
deepmd/pd/model/descriptor/se_t_tebd.py
[notice] 386-386: Unused local variable
Variable env_mat is not used.
[notice] 451-451: Unused local variable
Variable nall is not used.
[notice] 770-770: Unused local variable
Variable nall is not used.
deepmd/pd/utils/env.py
[notice] 72-72: Unused global variable
The global variable 'inter_nthreads' is not used.
[notice] 72-72: Unused global variable
The global variable 'intra_nthreads' is not used.
[notice] 75-76: Commented-out code
This comment appears to contain commented-out code.
deepmd/pd/utils/utils.py
[notice] 55-55: Statement has no effect
This statement has no effect.
deepmd/pd/utils/nlist.py
[notice] 395-395: Unused local variable
Variable nall is not used.
deepmd/pd/model/descriptor/dpa2.py
[warning] 663-663: Variable defined multiple times
This assignment to 'env_mat' is unnecessary as it is redefined before this value is used.
This assignment to 'env_mat' is unnecessary as it is redefined before this value is used.
[warning] 678-678: Variable defined multiple times
This assignment to 'env_mat' is unnecessary as it is redefined before this value is used.
[notice] 694-694: Unused local variable
Variable env_mat is not used.
deepmd/pd/train/training.py
[warning] 415-415: Unreachable code
This statement is unreachable.
[warning] 647-651: File is not always closed
File is opened but is not closed.
[warning] 657-657: File is not always closed
File is opened but is not closed.
[notice] 875-875: Unused local variable
Variable module is not used.
deepmd/pd/model/descriptor/repformer_layer.py
[notice] 967-967: Unused local variable
Variable ng2 is not used.
[notice] 1277-1277: Unused local variable
Variable nitem is not used.
deepmd/pd/model/descriptor/repformers.py
[notice] 416-416: Unused local variable
Variable ng1 is not used.
deepmd/pd/model/descriptor/se_atten.py
[notice] 431-431: Unused local variable
Variable nall is not used.
deepmd/pd/model/descriptor/dpa1.py
[notice] 529-529: Unused local variable
Variable env_mat is not used.
[notice] 600-600: Unused local variable
Variable nall is not used.
deepmd/pd/model/descriptor/se_a.py
[notice] 50-50: Empty except
'except' clause does nothing but pass and there is no explanatory comment.
[notice] 323-323: Unused local variable
Variable env_mat is not used.
[warning] 655-655: Variable defined multiple times
This assignment to 'ti' is unnecessary as it is redefined before this value is used.
🔇 Additional comments (104)
deepmd/infer/deep_eval.py (1)
506-506: LGTM: Improved error message formatting
The error message for fparam size mismatch has been updated to use f-string formatting, which is more readable and maintainable.
deepmd/jax/env.py (1)
21-21: LGTM! Verify usage of jax_export.
The addition of jax_export to __all__ is consistent with the import statement. Let's verify where this newly exposed functionality is being used.
✅ Verification successful
LGTM! The exposed jax_export functionality is being used appropriately.
The verification shows that jax_export is actively used in two key modules:
deepmd/jax/model/hlo.py: For deserializing StableHLO modelsdeepmd/jax/utils/serialization.py: For exporting JIT-compiled functions and handling symbolic shapes
The addition to __all__ properly aligns with these existing usages, making the public interface consistent with the codebase's requirements.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Find where jax_export is being imported/used in the codebase
# This helps understand the impact of exposing this functionality
# Search for imports of jax_export through this module
rg -l "from deepmd.jax.env import.*jax_export"
# Search for direct usage of jax_export
rg "jax_export\."
Length of output: 681
deepmd/pt/model/task/dos.py (2)
Line range hint 50-60: LGTM: Case embedding parameter addition
The new dim_case_embd parameter is properly integrated with appropriate type hints and default value.
Line range hint 77-85: LGTM: Parent class initialization
The dim_case_embd parameter is correctly passed to the parent class InvarFitting.
.pre-commit-config.yaml (2)
63-63: LGTM: Minor version bump for clang-format
The update from v19.1.3 to v19.1.4 is a minor version bump that likely contains bug fixes. This change is safe.
32-32: Verify compatibility with ruff v0.8.0
The update from v0.7.4 to v0.8.0 is a significant version bump that may introduce breaking changes. Please ensure compatibility with the codebase.
deepmd/pt/model/network/init.py (1)
133-134: LGTM! Improved type checking for leaky_relu param.
The enhanced type checking correctly handles edge cases by explicitly excluding boolean values while accepting integers and floats. This prevents potential issues where boolean values could be incorrectly interpreted as valid negative slopes.
Let's verify the usage of this function with leaky_relu across the codebase:
✅ Verification successful
Type safety improvement in calculate_gain is safe to use
The verification shows that all usages of calculate_gain in the codebase are through the kaiming_normal_ initialization, and they consistently use nonlinearity="linear". Since the type safety improvement in calculate_gain specifically addresses the leaky_relu case, and no instances of leaky_relu usage were found, the changes are safe and won't impact any existing code.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Find all usages of calculate_gain with leaky_relu to ensure compatibility
# with the new type checking logic
# Search for direct calls to calculate_gain with leaky_relu
rg -l "calculate_gain.*leaky_relu"
# Search for kaiming initialization calls that might use leaky_relu (default nonlinearity)
rg -l "kaiming_(uniform|normal)_"
Length of output: 233
Script:
#!/bin/bash
# Let's examine the actual usage of calculate_gain and kaiming initialization
# to ensure the type safety changes don't break existing code
# Check calculate_gain usage in mlp.py files
echo "=== MLP Usage in PyTorch implementation ==="
rg "calculate_gain|kaiming_" deepmd/pt/model/network/mlp.py -B 2 -A 2
echo -e "\n=== Network Usage in PyTorch implementation ==="
rg "calculate_gain|kaiming_" deepmd/pt/model/network/network.py -B 2 -A 2
echo -e "\n=== MLP Usage in PaddlePaddle implementation ==="
rg "calculate_gain|kaiming_" deepmd/pd/model/network/mlp.py -B 2 -A 2
Length of output: 1627
deepmd/dpmodel/atomic_model/__init__.py (1)
51-56: LGTM! The reordering improves readability.
The reordering of exported entities in __all__ maintains all functionality while grouping related models together in a more logical sequence. The energy models, polar models, and linear/ZBL models are now grouped together, followed by utility functions.
Let's verify if this ordering aligns with any existing conventions:
✅ Verification successful
Looking at the output, I can see that the ordering pattern in __all__ lists across the codebase follows a consistent approach where related entities are grouped together. For example:
- In
deepmd/pt/model/atomic_model/__init__.py, we see a similar grouping:
"BaseAtomicModel",
"DPAtomicModel",
"DPDOSAtomicModel",
"DPDipoleAtomicModel",
"DPEnergyAtomicModel",
"DPPolarAtomicModel",
"DPPropertyAtomicModel",
"DPZBLLinearEnergyAtomicModel",
"LinearEnergyAtomicModel",
"PairTabAtomicModel",- In
deepmd/tf/model/__init__.py, similar models are grouped:
"DOSModel",
"DipoleModel",
"EnerModel",
"GlobalPolarModel",
"LinearEnergyModel",
"PairTabModel",
"PolarModel",The reordering aligns well with codebase conventions.
The new ordering in the file under review follows the same pattern seen across the codebase where related models are grouped together (energy models, polar models, linear/ZBL models, etc.). This consistent grouping improves readability and maintainability.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for similar ordering patterns in other __init__.py files
# Look for other __init__.py files with __all__ declarations to verify consistency
rg -U "__all__\s*=\s*\[[\s\S]*?\]" --type py --glob "*__init__.py"
Length of output: 17492
deepmd/jax/model/__init__.py (1)
23-23: LGTM! Verify import statements in dependent modules.
The changes maintain consistency between imports and exports. The ordering is logical and all imported classes are properly included in __all__.
Let's verify that dependent modules are not affected by the reordering:
Also applies to: 25-25
deepmd/dpmodel/model/__init__.py (1)
32-32: Verify the impact of exposing DPModelCommon in the public API
Adding DPModelCommon to __all__ makes it part of the public API. Since this is a WIP PR, we should verify:
- The necessity of exposing this class publicly
- The existence of proper documentation for this class
- The potential impact on downstream code
Let's verify these aspects:
✅ Verification successful
Adding DPModelCommon to public API is appropriate and well-supported
The verification shows that:
DPModelCommonis a well-documented base class that implements common methods for all Models, with proper docstrings.- It's already extensively used across the codebase as a base class for various model implementations (EnergyModel, DipoleModel, DOSModel, etc.).
- The class is consistently implemented across different backends (pd, pt, dpmodel).
- Existing imports show it's already being used in tests and model implementations.
Making it part of the public API is justified as it:
- Provides common functionality needed by model implementations
- Is already stable and well-integrated into the codebase
- Has proper documentation and type hints
- Follows consistent patterns across different backend implementations
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check documentation and usage of DPModelCommon
# Check if class has proper docstrings
echo "Checking DPModelCommon documentation..."
ast-grep --pattern 'class DPModelCommon {
$$$
}'
# Check for existing usage patterns
echo "Checking current usage patterns..."
rg -l "DPModelCommon" --type python
# Check if this might break existing imports
echo "Checking potential import conflicts..."
rg -l "from deepmd.dpmodel.model import \*" --type python
Length of output: 465
Script:
#!/bin/bash
# Let's try a different approach without file type filters
# Check class definition and documentation
echo "Checking DPModelCommon class definition..."
rg -A 10 "class DPModelCommon"
# Check current usage patterns
echo "Checking current usage patterns..."
rg "DPModelCommon"
# Check potential import conflicts
echo "Checking import patterns..."
rg "from deepmd.dpmodel.model import"
Length of output: 6582
deepmd/dpmodel/fitting/property_fitting.py (2)
Line range hint 81-103: Implementation aligns with DPA-2 support objectives.
The changes to support DPA-2 in the Paddle backend are well-structured:
- The new
dim_case_embdparameter enables case-specific embeddings - Version bump properly signals the API change
- Implementation maintains type safety and backward compatibility through default values
Also applies to: 116-116
116-116: Verify version compatibility implications.
The version compatibility check has been updated from 2 to 3. This change could affect backward compatibility with existing models.
Let's check for version-related documentation and other version checks:
✅ Verification successful
The search results reveal a clear pattern: version 3 checks are consistently used in fitting-related modules across different implementations (tf, pt, dpmodel) and specifically in DPA2-related code. This version bump appears to be part of a deliberate architectural change related to the DPA-2 support mentioned in the PR objectives.
Version compatibility changes are consistent with DPA-2 implementation
The version bump from 2 to 3 in property_fitting.py aligns with other fitting modules and DPA2 descriptors:
- All fitting modules (
ener_fitting.py,property_fitting.py,invar_fitting.py, etc.) use version 3 - DPA2-specific modules (
dpa2.py) explicitly require version 3 - Legacy descriptors and models maintain version 2 compatibility
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for version-related documentation and consistency
# Check for version numbers in documentation
echo "Checking for version documentation..."
rg -i "version.*3" docs/ || echo "No version 3 documentation found"
# Check for other version checks to ensure consistency
echo "Checking for other version checks..."
rg "check_version_compatibility.*2|check_version_compatibility.*3" || echo "No other version checks found"
Length of output: 5919
deepmd/env.py (1)
17-21: LGTM! Verify usage of newly exported variables.
The addition of precision-related variables to __all__ looks good and maintains alphabetical order. These exports are likely needed for the Paddle backend implementation.
Let's verify the usage of these newly exported variables:
✅ Verification successful
Exports are correctly used across the codebase
The newly exported variables from deepmd/env.py are extensively used throughout the codebase:
GLOBAL_ENER_FLOAT_PRECISIONandGLOBAL_NP_FLOAT_PRECISIONare used for precision control in TensorFlow, PyTorch, and JAX backendsSHARED_LIB_MODULEis used for shared library configurationglobal_float_precis used in summary reporting
The exports are necessary as these variables are fundamental to the framework's precision control and are required by multiple components including the new Paddle backend.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for usage of the newly exported variables in the Paddle backend
rg "GLOBAL_ENER_FLOAT_PRECISION|GLOBAL_NP_FLOAT_PRECISION|global_float_prec" --type python
Length of output: 123
Script:
#!/bin/bash
# Search for usage of the newly exported variables without file type restriction
rg "GLOBAL_ENER_FLOAT_PRECISION|GLOBAL_NP_FLOAT_PRECISION|SHARED_LIB_MODULE|global_float_prec"
Length of output: 38234
deepmd/pt/model/model/make_model.py (1)
517-519: LGTM! Verify atomic model implementation.
The set_case_embd method correctly delegates to the atomic model. The implementation is clean and follows the established pattern.
Let's verify that the atomic model implements this method:
✅ Verification successful
Method implementation verified across atomic models
The set_case_embd method is properly implemented across all atomic model types:
DPAtomicModelin both PT and PD implementationsPairTabAtomicModelLinearAtomicModeland its variants- Base atomic model interface
All implementations follow the same pattern of setting case embedding and properly delegate to their fitting networks when needed. The method under review correctly follows this established pattern.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for set_case_embd implementation in atomic models
ast-grep --pattern 'def set_case_embd(self, case_idx: int):'
Length of output: 62
Script:
#!/bin/bash
# Let's try a broader search for set_case_embd method using ripgrep
rg "set_case_embd" -A 3
Length of output: 9408
deepmd/pt/model/atomic_model/dp_atomic_model.py (1)
96-101: LGTM! Clean implementation of case embedding support.
The implementation is straightforward and follows good practices:
- Clear documentation
- Single responsibility principle
- Proper delegation to the fitting net
deepmd/dpmodel/fitting/general_fitting.py (5)
108-108: LGTM! Clean addition of case embedding dimension parameter.
The dim_case_embd parameter is properly added with a sensible default value of 0.
Also applies to: 131-131
176-179: LGTM! Proper initialization of case embedding.
The case embedding is correctly initialized based on the dimension parameter, with proper type handling.
232-237: LGTM! Well-implemented case embedding method.
The set_case_embd method correctly uses numpy's eye matrix to create one-hot encodings for case embeddings.
446-458: LGTM! Proper integration of case embedding in forward pass.
The case embedding is correctly concatenated with the input tensor and properly handled for both regular and zero inputs.
308-308: Version bump and serialization updates are properly handled.
The version is correctly bumped from 2 to 3, and case embedding is properly included in serialization.
Also applies to: 316-316, 325-325
deepmd/pt/model/atomic_model/linear_atomic_model.py (2)
161-167: LGTM! Well-implemented case embedding for linear energy model.
The implementation correctly propagates the case embedding to all constituent models.
572-578: LGTM! Proper selective case embedding for ZBL model.
The implementation correctly sets the case embedding only for the DP model (first model), which is the expected behavior as ZBL models don't use case embeddings.
deepmd/dpmodel/fitting/__init__.py (2)
25-31: LGTM! The reordering improves readability.
The reordering of items in __all__ follows a more logical alphabetical order, which enhances code maintainability.
Line range hint 1-31: Verify imports match exports.
Let's ensure all imported entities are properly exported and vice versa.
✅ Verification successful
The previous script didn't correctly capture the imports due to the pattern matching. Let me try a different approach to extract the imports.
All imported entities are properly exported in __all__
After verifying the imports and exports, I can confirm that there's a perfect match between imported and exported entities:
Imported:
- DipoleFitting
- DOSFittingNet
- EnergyFittingNet
- InvarFitting
- make_base_fitting
- PolarFitting
- PropertyFittingNet
All of these are correctly listed in the __all__ list in the same order. There are no missing exports or extra imports.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify consistency between imports and exports in __init__.py
# Extract imported names and __all__ list contents
echo "Checking deepmd/dpmodel/fitting/__init__.py for import/export consistency..."
# Get all imported names (excluding 'from' and 'import' keywords)
imported=$(awk -F'[{,}]' '/^from.*import.*/ {
# Extract everything between "import (" and ")"
match($0, /import \((.*)\)/, arr)
if (arr[1] != "") {
# Clean up the captured group
gsub(/[ \n]/, "", arr[1])
print arr[1]
}
}' deepmd/dpmodel/fitting/__init__.py)
# Get all exported names from __all__
exported=$(awk '/^__all__/,/]/ {
# Remove __all__ = [ and ] and clean up
gsub(/^__all__ = \[|]|"| /, "")
if (NF > 0) print
}' deepmd/dpmodel/fitting/__init__.py)
# Compare the lists
echo "Imported names:"
echo "$imported"
echo "Exported names:"
echo "$exported"
# Check if any names are in one list but not the other
echo "Checking for mismatches..."
Length of output: 1151
Script:
#!/bin/bash
# Description: Verify consistency between imports and exports using grep
echo "Imported names:"
grep -o "from .* import.*(" deepmd/dpmodel/fitting/__init__.py | while read -r line; do
# Get the next line after each import statement that contains a closing parenthesis
line_num=$(grep -n "$line" deepmd/dpmodel/fitting/__init__.py | cut -d: -f1)
sed -n "${line_num},/)/p" deepmd/dpmodel/fitting/__init__.py
done
echo -e "\nExported names:"
grep -A 10 "^__all__" deepmd/dpmodel/fitting/__init__.py
Length of output: 2004
deepmd/dpmodel/fitting/invar_fitting.py (3)
126-126: LGTM: Case embedding dimension parameter added correctly.
The new parameter dim_case_embd follows the established pattern and maintains backward compatibility with its default value of 0.
188-188: Verify version bump documentation.
The version compatibility check has been updated from 2 to 3. Please ensure this breaking change is documented in the changelog.
159-159: Verify superclass handling of case embedding dimension.
The parameter is correctly passed to the superclass. Please ensure that the superclass GeneralFitting properly handles this new parameter.
✅ Verification successful
Let me check the constructor and parameter handling in GeneralFitting class.
Let me check how the dim_case_embd parameter is used within the GeneralFitting class.
Parameter is properly handled by the superclass
The GeneralFitting class correctly handles the dim_case_embd parameter:
- It's defined in the constructor with a default value of 0
- The value is stored in an instance variable
- It's used conditionally when
dim_case_embd > 0for case embedding initialization - The parameter is included in the class's state management
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if GeneralFitting class properly handles dim_case_embd
# Look for the parameter in the parent class implementation
rg -A 10 "class GeneralFitting" deepmd/dpmodel/fitting/
Length of output: 800
Script:
#!/bin/bash
# Check the constructor and parameter handling in GeneralFitting
rg -B 5 -A 20 "__init__" deepmd/dpmodel/fitting/general_fitting.py
Length of output: 977
Script:
#!/bin/bash
# Search for usage of dim_case_embd in GeneralFitting implementation
rg "dim_case_embd" deepmd/dpmodel/fitting/general_fitting.py
Length of output: 441
deepmd/pt/model/task/dipole.py (3)
86-86: LGTM! Parameter addition is well-structured
The new parameter follows the class's parameter organization pattern and maintains backward compatibility with its default value.
109-109: LGTM! Proper parent class initialization
The new parameter is correctly passed to the parent class.
135-135: Verify version compatibility handling
The version check has been updated from 2 to 3, which suggests a breaking change in the serialization format. Please ensure that:
- This change is documented in migration guides
- Other components that deserialize this class are updated accordingly
deepmd/pt/cxx_op.py (1)
53-55: LGTM! The f-string formatting improves readability.
The conversion to f-strings makes the error message more maintainable and easier to read while maintaining the same functionality.
Let's verify the error message formatting:
✅ Verification successful
The f-string formatting is correctly implemented and properly handles error messages
After examining the complete file context, the f-string formatting in the error messages is well-structured and properly implemented:
- All f-strings have properly closed brackets
- Variables are correctly interpolated
- The multi-line error messages maintain readability
- The error messages provide clear information about ABI and version incompatibilities
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify the error message formatting by checking for any string formatting issues
# or unclosed f-string brackets in the codebase
# Check for potential f-string formatting issues
rg -U "f['\"].*?['\"]" deepmd/pt/cxx_op.py | while read -r line; do
if echo "$line" | grep -q "{.*{\\|}.*}"; then
echo "Potential nested brackets in f-string: $line"
fi
if echo "$line" | grep -q "{[^}]*$"; then
echo "Potential unclosed bracket in f-string: $line"
fi
done
# Check for consistency in error message formatting
ast-grep --pattern 'raise RuntimeError(f"$_")' deepmd/pt/cxx_op.py
Length of output: 2072
Script:
#!/bin/bash
# Let's check the complete context of the error message to ensure proper formatting
rg -U "raise RuntimeError.*?\)" deepmd/pt/cxx_op.py -A 5
Length of output: 202
Script:
#!/bin/bash
# Let's get more context around the error message construction and usage
rg -U "error_message.*?=.*?\n.*?raise RuntimeError" deepmd/pt/cxx_op.py -B 10
Length of output: 79
Script:
#!/bin/bash
# Let's examine the entire file content to understand the error message context
cat deepmd/pt/cxx_op.py
Length of output: 3424
deepmd/pd/utils/exclude_mask.py (2)
29-61: LGTM! Well-implemented type masking logic
The implementation is clean and efficient:
- Properly converts types to mask using numpy array
- Correctly reshapes tensor for broadcasting
- Clear and complete documentation
80-156: LGTM! Efficient implementation with good handling of edge cases
The implementation is well thought out:
- Properly handles symmetric pairs by adding both (i,j) and (j,i)
- Efficiently handles virtual atoms for neighbor list masks
- Uses optimized tensor operations for computing pair masks
- Includes early return optimization for no exclusions case
deepmd/pt/model/model/__init__.py (3)
208-208: LGTM: Consistent model parameter storage
Adding model_def_script ensures ZBL models store their parameters consistently with other model types, which is good for reproducibility and debugging.
300-312: LGTM: Improved organization of exported symbols
The reordering of __all__ improves readability by grouping related models together, with base models first, specific model types in the middle, and factory functions at the end.
Line range hint 1-312: Verify the relationship with DPA-2 Paddle support
While the changes in this PyTorch backend file are sound, they seem tangential to the PR's main objective of implementing DPA-2 support in the Paddle backend. Let's verify if there are corresponding changes in the Paddle backend.
✅ Verification successful
The changes in PyTorch backend are consistent with DPA-2 implementation in Paddle
The changes in the PyTorch backend file are justified as they maintain consistency with the DPA-2 implementation in the Paddle backend. The evidence shows:
- DPA-2 is properly implemented in Paddle backend at
deepmd/pd/model/descriptor/dpa2.py - DPA-2 is registered in Paddle's descriptor registry at
deepmd/pd/model/descriptor/__init__.py - The PyTorch backend changes ensure both backends have consistent model initialization and handling
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for corresponding changes in Paddle backend
# Look for similar files and DPA-2 related changes in the Paddle backend
# Check for corresponding model initialization files in Paddle backend
fd -t f "__init__.py" "deepmd/pd"
# Look for DPA-2 related files or changes
rg -l "dpa2|DPA-2|DPA2" "deepmd/pd"
Length of output: 519
Script:
#!/bin/bash
# Let's examine the corresponding model/__init__.py in Paddle backend and DPA2 changes
# Check the corresponding model/__init__.py in Paddle backend
cat deepmd/pd/model/model/__init__.py
# Look at DPA2 descriptor implementation
cat deepmd/pd/model/descriptor/dpa2.py
# Check descriptor/__init__.py for DPA2 registration
cat deepmd/pd/model/descriptor/__init__.py
Length of output: 40602
deepmd/pd/model/descriptor/repformers.py (2)
375-378: Well-implemented error handling and input validation
The code demonstrates good practices in error handling:
- Proper assertion of required parameters
- Clear error messages for missing statistics
- Protection against division by zero
Also applies to: 564-568
91-176: Excellent documentation and type hints
The code demonstrates exemplary documentation practices:
- Comprehensive docstrings with detailed parameter descriptions
- Proper type hints that enhance code readability and maintainability
- Clear method-level documentation
deepmd/dpmodel/fitting/dipole_fitting.py (2)
134-134: LGTM: Parameter correctly passed to superclass
The new parameter is properly forwarded to the superclass constructor, maintaining consistent parameter ordering.
164-164: Verify migration path for serialization version bump
The version compatibility check has been updated from 2 to 3. Please ensure there's a migration path for models serialized with older versions.
✅ Verification successful
Version bump is backward compatible
The version compatibility check has been updated from 2 to 3 in the dipole fitting module, but this is safe since the version check function accepts any version between 1 (minimal_supported_version) and 3 (maximum_supported_version). This means older serialized models with versions 1 and 2 will continue to work with the new code.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for version migration logic or documentation
rg -A 5 "check_version_compatibility.*2.*3" --type py
Length of output: 55
Script:
#!/bin/bash
# Let's first check for any version compatibility checks in the codebase
rg -A 5 "check_version_compatibility" --type py
# Also search for version-related constants or variables
rg "VERSION|version" --type py
Length of output: 65617
Script:
#!/bin/bash
# Let's check the version.py file to understand the version compatibility logic
cat deepmd/utils/version.py
# Also check for any version migration documentation
rg -l "version.*migration|migration.*version" --type md
Length of output: 1030
deepmd/pd/model/network/init.py (2)
1-456: Implementation looks solid!
The tensor initialization module is well-implemented with comprehensive documentation and proper error handling. The code successfully adapts PyTorch's initialization patterns to the Paddle framework while maintaining clarity and correctness.
🧰 Tools
🪛 Ruff (0.8.0)
56-56: Ambiguous variable name: l
(E741)
328-328: No explicit stacklevel keyword argument found
(B028)
378-378: No explicit stacklevel keyword argument found
(B028)
157-176: Verify tensor shape handling for different frameworks
The fan calculation assumes a specific tensor layout (channels first/last). Since this is adapted from PyTorch to Paddle, we should verify that the tensor shape assumptions remain valid.
✅ Verification successful
Let me gather more information about tensor shape handling by searching for specific patterns.
Let me check one more thing to verify the tensor layout conventions in Paddle.
Tensor shape handling is correctly implemented for Paddle framework
The implementation correctly handles tensor shapes with proper consideration for the Paddle framework:
- The
reverseparameter is explicitly documented to handle tensor data format order[fout, fin, ...]by default (whenreverse=False) - The shape handling is flexible and supports both format orders through the
reverseparameter - All initialization functions (kaiming_normal_, xavier_uniform_, etc.) properly propagate the
reverseparameter - The implementation includes proper validation for tensor dimensions and empty tensors
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for tensor shape assumptions in the codebase
rg -A 2 "tensor\.shape" --type python
# Search for Paddle tensor creation patterns
ast-grep --pattern 'paddle.empty($$$)'
Length of output: 427
Script:
#!/bin/bash
# Search for tensor shape handling in initialization functions
rg -A 5 "\.shape" deepmd/pd/model/network/init.py
# Search for tensor operations to understand layout conventions
rg -A 3 "tensor\." deepmd/pd/model/network/init.py
# Look for any framework-specific imports or configurations
rg -l "import paddle" deepmd/pd/model/network/
Length of output: 3092
Script:
#!/bin/bash
# Look for tensor layout documentation or usage in tests
rg -B 2 -A 2 "reverse" deepmd/pd/model/network/init.py
# Check if there are any test files that verify initialization
fd -e py test_init deepmd/pd/model/network/
Length of output: 3457
deepmd/pt/model/descriptor/dpa2.py (1)
407-409: Verify the removal of sharing levels 2 and 3.
The simplification of parameter sharing levels improves code maintainability. However, we should verify that no existing code relies on the removed sharing levels 2 and 3.
✅ Verification successful
Removal of sharing levels 2 and 3 is safe
Based on the comprehensive code search results, the codebase only uses sharing levels 0 and 1. The removed levels 2 and 3 were never implemented or used in any descriptor class, test files, or configuration examples. All descriptors (DPA1, DPA2, SE-A, SE-T, etc.) consistently implement only levels 0 (share all) and 1 (share type embedding).
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Search for usage of sharing levels 2 and 3 in the codebase
# Search for direct usage of share_params with levels 2 or 3
rg -A 5 'share_params.*(?:2|3)'
# Search for configuration files that might specify these sharing levels
rg -A 5 'shared_level.*(?:2|3)'
Length of output: 67
Script:
#!/bin/bash
# Let's try a broader search to find any references to shared_level or parameter sharing
# Search for shared_level in Python files
rg "shared_level" --type py -B 2 -A 2
# Search for share_params method calls or definitions
rg "share_params" --type py -B 2 -A 2
# Search for any test files that might be testing parameter sharing
fd "test.*\.py" --exec rg -l "share"
Length of output: 56070
Script:
#!/bin/bash
# Let's check the test files for any usage of sharing levels 2 and 3
rg -A 5 "shared_level.*[23]" ./source/tests/pd/test_multitask.py ./source/tests/pt/test_multitask.py
# Also check configuration examples for any usage of these levels
rg -A 5 "shared_level.*[23]" ./source/tests/common/test_examples.py
Length of output: 172
deepmd/dpmodel/fitting/polarizability_fitting.py (2)
104-104: LGTM: Case embedding dimension parameter added correctly
The addition of dim_case_embd parameter with a default value of 0 maintains backward compatibility while enabling case embedding support for DPA-2.
154-154: LGTM: Case embedding dimension properly propagated to base class
The dim_case_embd parameter is correctly passed to the superclass constructor, ensuring proper initialization of case embedding support.
deepmd/jax/descriptor/__init__.py (2)
Line range hint 1-35: LGTM: Module exports are well-organized.
The module's public interface is well-organized with:
- Clear imports grouped by descriptor type
- Consistent
__all__declaration matching the imports - Proper SPDX license header
29-29: Verify DPA-2 implementation completeness.
Since this PR's main objective is to implement DPA-2 support, let's verify the implementation.
✅ Verification successful
Let me gather more information about the DPA2 implementation by checking the actual class implementations and tests.
Let me check the test files to verify the implementation completeness.
Based on the comprehensive test files and implementations shown, I can now generate the final response:
DPA-2 implementation is complete and well-tested across frameworks
The DPA-2 descriptor has been fully implemented with:
- Complete implementations in JAX, PyTorch, and PaddlePaddle frameworks
- Comprehensive test coverage including:
- Common tests for basic functionality and self-consistency
- Cross-framework consistency tests
- Framework-specific tests with model compression
- Numerical accuracy tests with reference values
- Support for all key features:
- Three-body interactions
- Type embedding
- Repformer layers
- Model compression
- Various precision modes
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check the DPA-2 implementation across the codebase
# Check for DPA-2 related files and their contents
echo "Searching for DPA-2 related files..."
fd "dpa2" --type f
# Check for DPA-2 class implementation details
echo -e "\nChecking DPA-2 class implementation..."
ast-grep --pattern 'class DescrptDPA2' -A 10
# Look for DPA-2 tests
echo -e "\nSearching for DPA-2 tests..."
rg "test.*dpa2" --type py
Length of output: 1323
Script:
#!/bin/bash
# Check the implementation of DescrptDPA2 in different frameworks
echo "JAX Implementation:"
cat deepmd/jax/descriptor/dpa2.py
echo -e "\nPyTorch Implementation:"
cat deepmd/pt/model/descriptor/dpa2.py
echo -e "\nPaddlePaddle Implementation:"
cat deepmd/pd/model/descriptor/dpa2.py
echo -e "\nTest Implementation Details:"
cat source/tests/common/dpmodel/test_descriptor_dpa2.py
cat source/tests/consistent/descriptor/test_dpa2.py
Length of output: 65978
Script:
#!/bin/bash
# Check test files for DPA2
echo "Common Tests:"
cat source/tests/common/dpmodel/test_descriptor_dpa2.py
echo -e "\nConsistency Tests:"
cat source/tests/consistent/descriptor/test_dpa2.py
echo -e "\nPyTorch Tests:"
cat source/tests/pt/model/test_descriptor_dpa2.py
cat source/tests/pt/model/test_compressed_descriptor_dpa2.py
echo -e "\nPaddlePaddle Tests:"
cat source/tests/pd/model/test_descriptor_dpa2.py
Length of output: 37293
deepmd/infer/model_devi.py (1)
207-216: LGTM: Improved string formatting
The update from %-formatting to f-strings improves code readability while maintaining the same functionality and alignment.
deepmd/pd/utils/decomp.py (2)
70-92: LGTM! Well-implemented utility function
The function is well-documented, has proper input validation, and handles the edge case of non-divisible lengths correctly.
1-139: Verify usage patterns in DPA-2 implementation
Since this file provides critical tensor operations for DPA-2 support, let's verify its usage patterns:
✅ Verification successful
Implementation verified and properly used across backends
The verification shows that:
- The decomposed operations are correctly used in both Paddle and PyTorch backends with consistent patterns
scatter_reduceis used identically in bothtransform_output.pyimplementations for force and virial calculations- The operations are properly tested in
test_decomp.pyandtest_forward_lower.py - The implementation maintains consistency between backends while providing the necessary DPA-2 support
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for usage of these operations in DPA-2 related code
echo "Searching for usage of decomposed operations..."
rg -l "masked_add_|scatter_reduce" --type py
# Check if there are any similar operations in torch backend for comparison
echo "Checking torch backend implementation..."
rg -l "masked_add|scatter_reduce" --type py | grep "torch"
Length of output: 576
Script:
#!/bin/bash
# Let's examine the actual usage patterns in the identified files
echo "Checking usage in transform_output.py files..."
rg -B2 -A2 "masked_add_|scatter_reduce" deepmd/pd/model/model/transform_output.py deepmd/pt/model/model/transform_output.py
echo "\nChecking test implementations..."
rg -B2 -A2 "masked_add_|scatter_reduce" source/tests/pd/test_decomp.py source/tests/pd/model/test_forward_lower.py
# Check if these operations are used consistently across backends
echo "\nComparing implementations between paddle and torch..."
ast-grep --pattern 'def masked_add_($$$)'
ast-grep --pattern 'def scatter_reduce($$$)'
Length of output: 3303
deepmd/dpmodel/atomic_model/linear_atomic_model.py (2)
137-144: Add return type hint and verify sub-models support.
The implementation looks good, but consider these improvements:
- Add return type hint:
def set_case_embd(self, case_idx: int) -> None: - We should verify that all sub-models support the
set_case_embdmethod.
✅ Verification successful
Based on the search results, I can now generate a final response as we have sufficient information. The results show that:
- The
set_case_embdmethod is widely implemented across various model classes in different backends (pd, pt, dpmodel) - The models used in
LinearEnergyAtomicModelare instances of atomic models, which all implement theset_case_embdmethod as evidenced by its presence indp_atomic_model.py,pairtab_atomic_model.py, and other atomic model implementations - There's even a base implementation in
make_base_atomic_model.pythat includes the return type hint
Implementation is safe and follows established pattern
The set_case_embd method implementation is correct and safe to use as all potential sub-models in the codebase implement this method. The suggestion about adding a return type hint is still valid but not critical, as it follows the base implementation in make_base_atomic_model.py.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if all potential sub-model types support set_case_embd
# Look for class definitions that could be used as sub-models
ast-grep --pattern 'class $name($_) {
$$$
set_case_embd($$$) {
$$$
}
$$$
}'
Length of output: 89
Script:
#!/bin/bash
# Let's try a different approach to find all classes that might be used as sub-models
# and check their method definitions
# First, let's find where these models are instantiated or imported
rg -A 5 "self\.models\s*="
# Also search for set_case_embd method implementations
rg "def set_case_embd"
Length of output: 3947
439-446: 🛠️ Refactor suggestion
Improve robustness and type hints for case embedding.
The current implementation has several areas for improvement:
- The assumption that
models[0]is the DP model should be validated - Missing return type hint
- No validation for empty models list
Consider this safer implementation:
- def set_case_embd(self, case_idx: int):
+ def set_case_embd(self, case_idx: int) -> None:
"""
Set the case embedding of this atomic model by the given case_idx,
typically concatenated with the output of the descriptor and fed into the fitting net.
"""
- # only set case_idx for dpmodel
- self.models[0].set_case_embd(case_idx)
+ if not self.models:
+ raise ValueError("No models available")
+ dp_model = self.models[0]
+ if not isinstance(dp_model, DPAtomicModel):
+ raise ValueError("First model must be a DPAtomicModel")
+ dp_model.set_case_embd(case_idx)deepmd/pt/model/descriptor/__init__.py (1)
46-58: Verify DPA-2 public interface completeness.
Let's ensure all necessary components for DPA-2 support are properly exposed through the public interface.
✅ Verification successful
Let me gather more information about the DPA-2 related components to ensure completeness.
Let me search for the actual DPA-2 implementation to verify its completeness.
Based on the extensive code review, I can now provide a final response regarding the DPA-2 public interface completeness.
DPA-2 descriptor interface is complete and properly exposed.
The implementation shows that:
-
The DPA-2 descriptor is properly exposed in
deepmd/pt/model/descriptor/__init__.pywith all necessary components:- Main descriptor class
DescrptDPA2 - Supporting configuration classes
RepformerArgsandRepinitArgs - All required utility functions and type embeddings
- Main descriptor class
-
The implementation is complete with:
- Full PyTorch and PaddlePaddle backend support
- Comprehensive test coverage
- Proper serialization/deserialization support
- Complete configuration options including type embedding, three-body interactions, and attention mechanisms
-
The descriptor is being actively used in tests and integrated with the rest of the codebase, as evidenced by the test files and import patterns.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for any potential DPA-2 related classes/functions that might need to be exposed
# Search for DPA-2 related classes/functions in the codebase
echo "Searching for DPA-2 related definitions..."
ast-grep --pattern 'class $name {
$$$
}' | rg -A 5 'DPA2|dpa2'
# Search for DPA-2 related imports in other files
echo "Checking usage of DPA-2 components..."
rg -l "from deepmd.pt.model.descriptor import" | xargs rg "DescrptDPA2"
Length of output: 4644
Script:
#!/bin/bash
# Check for DPA-2 related parameter classes and their usage
echo "Searching for DPA2 parameter classes..."
ast-grep --pattern 'class $name {
$$$
}' | rg -B 2 -A 5 'DescriptorParamDPA2'
# Check the actual implementation file for DPA2
echo "Checking DPA2 implementation imports..."
fd "descriptor_dpa2.py" --exec cat {}
# Check if there are any utility functions related to DPA2
echo "Checking for DPA2 related utility functions..."
rg -l "def.*dpa2" | xargs cat
Length of output: 85589
Script:
#!/bin/bash
# Search for DPA-2 descriptor implementation
echo "Searching for DPA-2 descriptor implementation..."
fd "dpa2.py" --exec cat {}
# Search for DPA-2 related imports and usage
echo "Checking DPA-2 imports and usage..."
rg "class DescrptDPA2"
Length of output: 125228
deepmd/common.py (3)
38-44: LGTM! Well-organized exports.
The reordering of exports in __all__ improves readability by grouping constants first, followed by functions.
Line range hint 48-57: Verify integration with Paddle backend.
The type system changes appear to support DPA-2 implementation. Let's verify the integration with Paddle backend:
Also applies to: 65-69
✅ Verification successful
Integration with Paddle backend is properly implemented
The activation functions defined in _ACTIVATION are fully supported in the Paddle backend through the implementation in deepmd/pd/utils/utils.py. All activation functions (relu, relu6, softplus, sigmoid, tanh, gelu, gelu_tf, none, linear) are properly mapped to their corresponding Paddle implementations using paddle.nn.functional and paddle.tanh().
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for Paddle-specific usage of these types
rg --type python "import paddle|from paddle" -B 2 -A 2
# Search for DPA-2 related code that might use these types
rg --type python "dpa2|DPA2|DPA-2" -B 2 -A 2
Length of output: 165
Script:
#!/bin/bash
# Search for Paddle-specific usage without file type restriction
rg "import paddle|from paddle" -B 2 -A 2
# Search for activation functions usage in paddle-related files
rg -g "*.py" "(relu|relu6|softplus|sigmoid|tanh|gelu|gelu_tf|none|linear)" "$(fd -g "*paddle*.py")"
# Search for DPA-2 related code
rg "dpa2|DPA2|DPA-2" -B 2 -A 2
# Look for paddle backend implementation files
fd -g "*paddle*.py"
Length of output: 93330
Script:
#!/bin/bash
# Search for activation function usage in Paddle backend
rg -g "*.py" "(relu|relu6|softplus|sigmoid|tanh|gelu|gelu_tf|none|linear)" "deepmd/pd/"
# Search for activation function implementations in Paddle backend
ast-grep --pattern $'def $_(self, $_, $$$) {
$$$
paddle.nn.functional.$activation($$$)
$$$
}' deepmd/pd/
Length of output: 7219
Line range hint 48-57: Add docstrings for type literals and verify activation functions.
While the type definitions are well-structured, consider:
- Adding docstrings to describe valid values and usage.
- Verifying that all required activation functions for DPA-2 are included.
Let's verify the activation functions usage:
✅ Verification successful
Based on the search results from the executed scripts, I can now provide a definitive response about the activation functions verification:
All required activation functions for DPA-2 are properly included and implemented.
The codebase shows comprehensive support for activation functions through:
- The
_ACTIVATIONtype literal includes all standard activation functions: relu, relu6, softplus, sigmoid, tanh, gelu, gelu_tf, none, and linear - Implementation is verified in multiple components:
- Core activation function class
ActivationFnin pd/utils/utils.py implements all listed functions - Consistent usage across descriptor and fitting networks
- Proper handling in both forward and gradient computations
- Core activation function class
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for activation function usage in the codebase
rg --type python "activation.*=|activation.*:" -B 2 -A 2
Length of output: 89
Script:
#!/bin/bash
# Search for activation function usage in the codebase without file type restriction
rg "activation.*=|activation.*:" -B 2 -A 2
Length of output: 65581
Script:
#!/bin/bash
# Search for imports or usage of activation functions
rg -i "activation|relu|softplus|sigmoid|tanh|gelu" -B 2 -A 2
Length of output: 74677
Script:
#!/bin/bash
# Find Python files that might contain activation-related code
fd -e py -x cat {} | rg -i "activation|relu|softplus|sigmoid|tanh|gelu" -B 2 -A 2
Length of output: 89766
Script:
#!/bin/bash
# Use ast-grep to find function definitions or class methods related to activation
ast-grep --pattern 'def $_($$$) {
$$$
activation
$$$
}'
Length of output: 63
deepmd/pd/train/wrapper.py (2)
1-20: LGTM: Well-structured imports and type definitions.
The file header is properly licensed, imports are well-organized, and type hints are appropriately used.
67-73: Verify implementation plan for parameter sharing.
The share_params method is crucial for DPA-2 support but is currently not implemented. Please clarify:
- Is this intentional for this PR stage?
- Are there dependencies on PR pd: support dpa1 #4414 mentioned in the PR description?
✅ Verification successful
Implementation plan for parameter sharing is appropriate
The unimplemented share_params in deepmd/pd/train/wrapper.py is intentional for this PR stage. The codebase shows:
- A complete parameter sharing implementation exists in the PyTorch (pt) version at
deepmd/pt/train/wrapper.py - The DPA-2 descriptor and other components already have their
share_paramsmethods implemented in both pd and pt versions - The pattern follows a hierarchical sharing structure where the wrapper coordinates sharing between different model components
This matches the PR's staged approach where PaddlePaddle implementation follows the established PyTorch patterns.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check if parameter sharing is implemented in other files or branches
rg -g '!wrapper.py' "share_params.*multitask" .
Length of output: 49
Script:
#!/bin/bash
# Search for any implementation or references to share_params
rg "share_params" .
# Also check for multitask related code
rg "multitask" .
# Look for any parameter sharing patterns in Python files
ast-grep --pattern 'def share_params($$$)'
Length of output: 19420
deepmd/pd/model/network/network.py (2)
1-34: LGTM! Well-organized imports and utility function.
The imports are logically grouped and the Tensor utility function provides a clean abstraction for tensor creation with consistent precision and device settings.
1-345: Overall implementation is solid and well-documented.
The code successfully implements type embedding networks for the Paddle backend with good separation of concerns and thorough documentation. While there are some areas for improvement in error handling and code organization, the implementation is robust and aligns well with the PR's objective of adding DPA-2 support.
deepmd/pd/model/task/invar_fitting.py (2)
35-37: Well-structured implementation with good practices
The InvarFitting class follows good practices with:
- Proper type hints and comprehensive documentation
- Clear separation of concerns in methods
- Appropriate use of decorators for registration and output checking
150-150: Verify version compatibility requirements
The version compatibility check uses hardcoded version numbers (3, 1). Consider:
- Documenting why these specific versions are required
- Using constants for version numbers
Let's check version compatibility usage across the codebase:
✅ Verification successful
Based on the analysis of the codebase, I can now provide a definitive response about the version compatibility check:
Version compatibility check is correctly implemented
The version numbers (3, 1) in check_version_compatibility(data.pop("@version", 1), 3, 1) are:
- 3: Maximum supported version for the invar_fitting model
- 1: Minimum supported version (default value)
This is part of a well-established versioning system across the codebase where:
- Each model component maintains its own version compatibility range
- The version check ensures backward compatibility while allowing for model evolution
- The version numbers are consistently used across different model components with clear boundaries
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for version compatibility checks
rg "check_version_compatibility" --type python -B 2 -A 2
Length of output: 87
Script:
#!/bin/bash
# Search for version compatibility checks without file type restriction
rg "check_version_compatibility" -B 2 -A 2
Length of output: 50187
Script:
#!/bin/bash
# Find the definition of check_version_compatibility
ast-grep --pattern 'def check_version_compatibility($$$)'
Length of output: 59
Script:
#!/bin/bash
# Find all Python files that might contain version-related constants
rg "VERSION|version" --type-add 'py:*.py' --type py
Length of output: 65918
Script:
#!/bin/bash
# Find the definition of check_version_compatibility in version.py
cat deepmd/utils/version.py
Length of output: 972
Script:
#!/bin/bash
# Find any documentation or comments about version compatibility in the codebase
rg -B 2 -A 2 "version.*compatibility|compatibility.*version" --type-add 'py:*.py' --type py
Length of output: 51040
deepmd/pd/model/model/__init__.py (3)
1-41: LGTM! Well-structured module setup with clear documentation.
The module documentation and import organization follow best practices, clearly explaining the purpose and relationships between components.
137-144: LGTM! Well-defined public API.
The exports clearly define the module's public interface, including all necessary components.
1-144: Verify integration with DPA-2 and Paddle backend.
Please ensure this implementation aligns with the DPA-2 requirements and integrates properly with the Paddle backend changes.
Let's verify the integration:
✅ Verification successful
Based on the verification results, I can now provide a final response:
Implementation properly integrates with DPA-2 and Paddle backend
The code shows correct integration with both DPA-2 and Paddle backend:
- Proper inheritance from Paddle's
nn.Layerin model implementations - Complete DPA-2 descriptor implementation with required components (repinit, repformer)
- Correct model factory pattern that supports both standard and custom model types
- Appropriate error handling for model configuration and initialization
The only potential issues that could arise are properly handled with clear error messages:
- Invalid fitting types
- Mismatched preset bias configurations
- Type map inconsistencies
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for DPA-2 and Paddle backend integration points
# Look for DPA-2 related configurations and usage
echo "Checking for DPA-2 configurations..."
rg -l "dpa2|DPA-2|DPA2" --type py
# Check for Paddle backend integration
echo "Checking for Paddle backend integration..."
rg -l "paddle|Paddle" --type py
# Look for related model implementations
echo "Checking for related model implementations..."
fd -e py -x rg -l "class.*Model.*paddle" {}
Length of output: 7069
Script:
#!/bin/bash
# Let's verify the model integration more specifically
# Check model implementation in pd backend
echo "=== Checking Paddle model implementation ==="
rg "class.*Model" deepmd/pd/model/model/model.py -A 10
# Check DPA2 descriptor implementation
echo -e "\n=== Checking DPA2 descriptor implementation ==="
rg "class" deepmd/pd/model/descriptor/dpa2.py -A 5
# Check model factory and integration points
echo -e "\n=== Checking model factory implementation ==="
rg "def get_model|def get_standard_model" deepmd/pd/model/model/__init__.py -A 10
# Check for any potential integration issues
echo -e "\n=== Checking for potential integration issues ==="
rg "RuntimeError|ValueError" deepmd/pd/model/model/__init__.py -A 2
Length of output: 5902
deepmd/pd/model/descriptor/se_atten.py (1)
990-1005: Handle Potential Division by Zero in Softmax
When computing attn_weights, there is a possibility of encountering -inf values when applying paddle.masked_fill with float("-inf"). This can lead to NaN values after applying softmax. Consider adding a small epsilon to prevent division by zero or handling this case explicitly.
Run the following script to check for any NaN values in attn_weights after softmax:
deepmd/dpmodel/atomic_model/dp_atomic_model.py (1)
68-73: LGTM! Clean implementation of case embedding support.
The implementation correctly delegates the case embedding to the fitting network and is well-documented.
deepmd/dpmodel/atomic_model/make_base_atomic_model.py (1)
71-77: LGTM! Well-defined abstract method for case embedding support.
The abstract method is properly defined with clear documentation and correct type annotations.
deepmd/dpmodel/atomic_model/pairtab_atomic_model.py (1)
123-130: LGTM! Clear handling of unsupported case embedding.
The implementation correctly raises NotImplementedError with a clear message.
deepmd/pd/model/task/ener.py (1)
75-75: Verify version compatibility change impact
The version compatibility check has been updated from 2 to 3. This change might affect backward compatibility with existing serialized models.
Run the following script to check for version compatibility documentation and migration guides:
✅ Verification successful
Let me gather more information about the version compatibility changes.
Let me gather more information about the version compatibility check implementation.
Let me check the version compatibility implementation.
Let me check one more thing about the version compatibility implementation.
Based on the gathered information, I can now provide the final response:
Version compatibility change is consistent with model evolution
The version compatibility check has been updated from 2 to 3 in deepmd/pd/model/task/ener.py as part of a coordinated update across multiple components:
- The change is consistent with version checks in other components (pd, tf, and pt backends)
- The version compatibility system properly handles backward compatibility through version ranges
- Migration utilities exist to convert models between versions
- The implementation includes proper error handling and clear error messages
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for version compatibility documentation
# Check README and docs for version compatibility info
rg -i "version.*compatibility|migration" README.md docs/
# Check for other version checks in the codebase
rg "check_version_compatibility.*3" --type py
Length of output: 1905
Script:
#!/bin/bash
# Search for version-related changes in git history
git log -n 10 --grep="version" --pretty=format:"%h %s"
# Look for version-related constants or utilities
rg "VERSION|version" --type py -B 2 -A 2
# Check for any test files related to version compatibility
fd "test.*version.*\.py$"
Length of output: 68283
Script:
#!/bin/bash
# Check for version compatibility check implementation
rg "check_version_compatibility" --type py -B 2 -A 2
# Look for version-related changes in the model
rg "MODEL_VERSION" --type py -B 2 -A 2
Length of output: 60402
Script:
#!/bin/bash
# Check the version compatibility implementation
rg -A 10 "def check_version_compatibility" deepmd/utils/version.py
# Look for version-related changes in git history
git log -n 5 --oneline deepmd/pd/model/task/ener.py
Length of output: 379
Script:
#!/bin/bash
# Get the complete version compatibility implementation
cat deepmd/utils/version.py
# Check for any version migration utilities
rg "convert.*version" --type py
Length of output: 1919
deepmd/dpmodel/descriptor/__init__.py (1)
31-33: Verify descriptor visibility requirements
The AI summary indicates that DescrptDPA1, DescrptDPA2, and DescrptHybrid should be removed from __all__, but they are still listed. Please clarify if these descriptors should remain public.
deepmd/dpmodel/utils/__init__.py (1)
39-61: Verify utility function visibility requirements
The AI summary indicates that several items should be removed from __all__, but they are still present:
EnvMatload_dp_model,save_dp_modeltraverse_model_dictbuild_neighbor_listnlist_distinguish_typesget_multiple_nlist_key
Please clarify if these utilities should remain public.
deepmd/pd/utils/region.py (3)
35-55: Implementation of inter2phys is correct
The function correctly converts internal coordinates to physical coordinates using matrix multiplication.
57-76: Implementation of to_face_distance is correct
The function accurately computes the distances to the faces of the simulation cell.
93-112: Implementation of normalize_coord is correct
The normalize_coord function correctly applies periodic boundary conditions to atomic coordinates.
deepmd/pt/model/task/fitting.py (10)
67-67: No issues with the updated comment.
The comment accurately reflects that bias_atom_e and case_embd are not shared.
98-99: Docstring updated to include dim_case_embd parameter.
The addition of dim_case_embd with its description is correctly documented.
137-137: Constructor updated to include dim_case_embd parameter.
The __init__ method now accepts dim_case_embd, enhancing flexibility.
159-159: Assigned self.dim_case_embd in __init__.
The instance variable dim_case_embd is properly initialized.
224-224: Adjusted input dimension to include dim_case_embd.
Input dimension in_dim now correctly accounts for dim_case_embd.
284-284: Updated serialization version to 3.
The version update in serialize reflects the new changes.
292-292: Included dim_case_embd in serialization data.
dim_case_embd is now part of the serialized output.
301-301: Serialized case_embd buffer.
The case_embd attribute is correctly added to serialized variables.
382-383: Extended __setitem__ to handle case_embd.
The __setitem__ method now supports setting case_embd.
400-401: Extended __getitem__ to handle case_embd.
The __getitem__ method now supports retrieving case_embd.
deepmd/dpmodel/fitting/ener_fitting.py (3)
35-35: Added dim_case_embd parameter to constructor.
EnergyFittingNet now accepts dim_case_embd, enhancing flexibility.
59-59: Passed dim_case_embd to superclass constructor.
dim_case_embd is correctly forwarded to the parent class.
78-78: Updated version compatibility check to version 3.
The version check in deserialize aligns with serialization updates.
deepmd/dpmodel/fitting/dos_fitting.py (3)
39-39: Added dim_case_embd parameter to constructor.
DOSFittingNet now includes dim_case_embd, enhancing model capabilities.
64-64: Passed dim_case_embd to superclass constructor.
dim_case_embd is appropriately passed to the parent class.
78-78: Updated version compatibility check to version 3.
Ensures compatibility with the new serialization format.
deepmd/pt/model/task/invar_fitting.py (4)
59-60: Docstring updated to include dim_case_embd parameter.
The documentation now reflects the additional parameter.
96-96: Added dim_case_embd parameter to constructor.
InvarFitting now accepts dim_case_embd, enhancing functionality.
119-119: Passed dim_case_embd to superclass constructor.
Parameter forwarding is correctly implemented.
148-148: Updated version compatibility check to version 3.
Alignment with updated serialization versions.
deepmd/pt/model/task/polarizability.py (5)
56-57: Docstring updated to include dim_case_embd parameter.
Documentation now includes the new parameter for clarity.
90-90: Added dim_case_embd parameter to constructor.
PolarFittingNet now supports dim_case_embd.
134-134: Passed dim_case_embd to superclass constructor.
Properly integrates dim_case_embd into the class hierarchy.
198-198: Updated serialization version to 4 in serialize method.
Reflects changes in serialization format.
209-209: Updated version compatibility check to version 4 in deserialize.
Ensures deserialization aligns with the updated serialization version.
deepmd/pd/utils/env.py (1)
28-31: Ensure consistent device selection and handle single-device environments.
The logic for setting DEVICE assumes that CUDA devices are available or sets it to CPU otherwise. However, in single-device environments or when LOCAL_RANK is not properly defined, DEVICE might not be set correctly. Ensure that DEVICE is correctly defined in all scenarios.
Run the following script to check the device assignment:
✅ Verification successful
Device selection logic is correctly implemented for both single and distributed environments.
The code analysis reveals that the device selection logic is properly implemented:
LOCAL_RANKis obtained frompaddle.distributed.get_rank(), which returns 0 for non-distributed settings- Device selection follows a clear logic:
- Sets to "cpu" if explicitly requested via environment variable or if no CUDA devices are available
- Otherwise sets to
gpu:{LOCAL_RANK}which correctly maps to GPU 0 in single-GPU environments
- The device setting is immediately applied via
paddle.device.set_device(DEVICE)
The implementation correctly handles all scenarios:
- Single-GPU environments (LOCAL_RANK = 0)
- Multi-GPU distributed environments (LOCAL_RANK assigned by distributed system)
- CPU-only environments
- Forced CPU usage through environment variable
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify that DEVICE is correctly assigned based on environment variables and CUDA availability.
python -c "
import os
import paddle
LOCAL_RANK = paddle.distributed.get_rank()
if os.environ.get('DEVICE') == 'cpu' or paddle.device.cuda.device_count() <= 0:
DEVICE = 'cpu'
else:
DEVICE = f'gpu:{LOCAL_RANK}'
print(f'DEVICE is set to: {DEVICE}')
"
Length of output: 384
Script:
#!/bin/bash
# Let's analyze the codebase to understand the device selection logic and its usage
# Check how LOCAL_RANK is defined and used
rg "LOCAL_RANK" -A 3 -B 3
# Check for any device-related configurations or settings
rg "DEVICE\s*=" -A 3 -B 3
# Look for paddle device-related code
rg "paddle.device" -A 3 -B 3
# Search for distributed training related code
rg "paddle.distributed" -A 3 -B 3
Length of output: 11454
deepmd/pd/model/descriptor/env_mat.py (1)
24-24: Correct Handling of Invalid Neighbor Indices with paddle.where
Using paddle.where to handle invalid neighbor indices ensures that computations do not break due to negative indices.
deepmd/entrypoints/test.py (1)
Line range hint 707-879: Verify file naming patterns across test functions.
Let's ensure that the file naming patterns are consistent across all test functions in the codebase.
✅ Verification successful
Based on the search results, I can now generate the final response:
File naming patterns are consistent across test functions
The file naming patterns follow a consistent convention throughout the test module:
- Single output files use simple suffixes:
.e.out,.f.out,.v.out - Per-atom variants append
_peratom:.e_peratom.out,.v_peratom.out - Iterative outputs use indexed suffixes:
.dos.out.{ii},.ados.out.{ii},.property.out.{ii},.aproperty.out.{ii}
All patterns use the with_suffix() method consistently for path manipulation, maintaining a uniform approach to file naming.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check file naming patterns in test functions
# Expected: Find all file output operations to verify naming consistency
# Search for file suffix patterns in test functions
rg -A 2 'with_suffix\(' deepmd/entrypoints/test.py
# Search for other file output operations that might use different patterns
rg -A 2 '\.out\.' deepmd/entrypoints/test.py
Length of output: 2799
deepmd/infer/__init__.py (1)
Line range hint 19-33: LGTM: Well-documented factory function
The DeepPotential factory function is well-documented with clear parameter descriptions and return type annotations.
deepmd/infer/deep_eval.py
Outdated
| @@ -516,8 +515,7 @@ def _standard_input(self, coords, cells, atom_types, fparam, aparam, mixed_type) | |||
| aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) | |||
| else: | |||
| raise RuntimeError( | |||
| "got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d" | |||
| % (nframes, natoms, fdim, natoms, fdim, fdim) | |||
| f"got wrong size of frame param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}" | |||
There was a problem hiding this comment.
Fix typo in error message
The error message incorrectly refers to "frame param" when it should be "atomic param" since this is checking the aparam dimensions.
Apply this diff to fix the typo:
- f"got wrong size of frame param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}"
+ f"got wrong size of atomic param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| f"got wrong size of frame param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}" | |
| f"got wrong size of atomic param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}" |
| rcut_smth: float, | ||
| sel: Union[list[int], int], | ||
| ntypes: int, | ||
| neuron: list = [2, 4, 8], |
There was a problem hiding this comment.
Fix mutable default arguments in function parameters
Using mutable objects (lists) as default arguments can lead to unexpected behavior because the same list object is shared across all function calls. Instead, use None as the default and initialize the list within the function.
Apply these changes:
- def __init__(self, ..., neuron: list = [2, 4, 8], ..., exclude_types: list[tuple[int, int]] = [], ...):
+ def __init__(self, ..., neuron: Optional[list] = None, ..., exclude_types: Optional[list[tuple[int, int]]] = None, ...):
super().__init__()
+ if neuron is None:
+ neuron = [2, 4, 8]
+ if exclude_types is None:
+ exclude_types = []Similar changes should be applied to the DescrptBlockSeTTebd class constructor as well.
Also applies to: 136-136, 521-521, 528-528, 733-733
🧰 Tools
🪛 Ruff (0.8.0)
129-129: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/utils/exclude_mask.py
Outdated
| def __init__( | ||
| self, | ||
| ntypes: int, | ||
| exclude_types: list[int] = [], | ||
| ) -> None: | ||
| super().__init__() | ||
| self.reinit(ntypes, exclude_types) |
There was a problem hiding this comment.
Fix mutable default arguments in method signatures
Using mutable default arguments in Python can lead to unexpected behavior when the same default list is shared across multiple instances.
Apply this fix:
def __init__(
self,
ntypes: int,
- exclude_types: list[int] = [],
+ exclude_types: list[int] | None = None,
) -> None:
super().__init__()
- self.reinit(ntypes, exclude_types)
+ self.reinit(ntypes, exclude_types or [])
def reinit(
self,
ntypes: int,
- exclude_types: list[int] = [],
+ exclude_types: list[int] | None = None,
) -> None:
+ exclude_types = exclude_types or []Also applies to: 22-26
🧰 Tools
🪛 Ruff (0.8.0)
17-17: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/utils/exclude_mask.py
Outdated
| def __init__( | ||
| self, | ||
| ntypes: int, | ||
| exclude_types: list[tuple[int, int]] = [], | ||
| ) -> None: | ||
| super().__init__() | ||
| self.reinit(ntypes, exclude_types) |
There was a problem hiding this comment.
Fix mutable default arguments in method signatures
Similar to AtomExcludeMask, using mutable default arguments should be avoided.
Apply this fix:
def __init__(
self,
ntypes: int,
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: list[tuple[int, int]] | None = None,
) -> None:
super().__init__()
- self.reinit(ntypes, exclude_types)
+ self.reinit(ntypes, exclude_types or [])
def reinit(
self,
ntypes: int,
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: list[tuple[int, int]] | None = None,
) -> None:
+ exclude_types = exclude_types or []Also applies to: 75-79
🧰 Tools
🪛 Ruff (0.8.0)
70-70: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
| update_residual_init: str = "norm", | ||
| set_davg_zero: bool = True, | ||
| smooth: bool = True, | ||
| exclude_types: list[tuple[int, int]] = [], |
There was a problem hiding this comment.
Replace mutable default arguments with None
Using mutable objects (lists) as default arguments is a Python pitfall as they are created once when the function is defined and shared between all calls. This can lead to unexpected behavior if the list is modified.
Apply this fix:
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: Optional[list[tuple[int, int]]] = None,Then initialize the list in the method body:
if exclude_types is None:
exclude_types = []Also applies to: 360-360
🧰 Tools
🪛 Ruff (0.8.0)
81-81: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/utils/utils.py
Outdated
| def to_numpy_array( | ||
| xx, | ||
| ): | ||
| if xx is None: | ||
| return None | ||
| assert xx is not None | ||
| # Create a reverse mapping of PD_PRECISION_DICT | ||
| reverse_precision_dict = {v: k for k, v in PD_PRECISION_DICT.items()} | ||
| # Use the reverse mapping to find keys with the desired value | ||
| prec = reverse_precision_dict.get(xx.dtype, None) | ||
| prec = NP_PRECISION_DICT.get(prec, np.float64) | ||
| if prec is None: | ||
| raise ValueError(f"unknown precision {xx.dtype}") | ||
| if isinstance(xx, np.ndarray): | ||
| return xx.astype(prec) | ||
| if xx.dtype == paddle.bfloat16: | ||
| xx = xx.astype(paddle.get_default_dtype()) | ||
| return xx.numpy().astype(prec) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove redundant assertion and handle type checking properly.
The function contains redundant checks and commented-out type checking code.
Apply these improvements:
def to_numpy_array(
xx,
):
if xx is None:
return None
- assert xx is not None
# Create a reverse mapping of PD_PRECISION_DICT
reverse_precision_dict = {v: k for k, v in PD_PRECISION_DICT.items()}
# Use the reverse mapping to find keys with the desired value
prec = reverse_precision_dict.get(xx.dtype, None)
prec = NP_PRECISION_DICT.get(prec, np.float64)
if prec is None:
raise ValueError(f"unknown precision {xx.dtype}")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def to_numpy_array( | |
| xx, | |
| ): | |
| if xx is None: | |
| return None | |
| assert xx is not None | |
| # Create a reverse mapping of PD_PRECISION_DICT | |
| reverse_precision_dict = {v: k for k, v in PD_PRECISION_DICT.items()} | |
| # Use the reverse mapping to find keys with the desired value | |
| prec = reverse_precision_dict.get(xx.dtype, None) | |
| prec = NP_PRECISION_DICT.get(prec, np.float64) | |
| if prec is None: | |
| raise ValueError(f"unknown precision {xx.dtype}") | |
| if isinstance(xx, np.ndarray): | |
| return xx.astype(prec) | |
| if xx.dtype == paddle.bfloat16: | |
| xx = xx.astype(paddle.get_default_dtype()) | |
| return xx.numpy().astype(prec) | |
| def to_numpy_array( | |
| xx, | |
| ): | |
| if xx is None: | |
| return None | |
| # Create a reverse mapping of PD_PRECISION_DICT | |
| reverse_precision_dict = {v: k for k, v in PD_PRECISION_DICT.items()} | |
| # Use the reverse mapping to find keys with the desired value | |
| prec = reverse_precision_dict.get(xx.dtype, None) | |
| prec = NP_PRECISION_DICT.get(prec, np.float64) | |
| if prec is None: | |
| raise ValueError(f"unknown precision {xx.dtype}") | |
| if isinstance(xx, np.ndarray): | |
| return xx.astype(prec) | |
| if xx.dtype == paddle.bfloat16: | |
| xx = xx.astype(paddle.get_default_dtype()) | |
| return xx.numpy().astype(prec) |
deepmd/pd/utils/utils.py
Outdated
| def get_generator( | ||
| seed: int | list[int] | None = None, | ||
| ) -> PaddleGenerator | None: | ||
| if seed is not None: | ||
| if isinstance(seed, list): | ||
| seed = mix_entropy(seed) | ||
| if DEVICE == "cpu": | ||
| generator = paddle.framework.core.default_cpu_generator() | ||
| elif DEVICE == "gpu": | ||
| generator = paddle.framework.core.default_cuda_generator(0) | ||
| elif DEVICE.startswith("gpu:"): | ||
| generator = paddle.framework.core.default_cuda_generator( | ||
| int(DEVICE.split("gpu:")[1]) | ||
| ) | ||
| else: | ||
| raise ValueError("DEVICE should be cpu or gpu or gpu:x") | ||
| generator.manual_seed(seed) | ||
| return generator | ||
| else: | ||
| return None |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Improve device validation and error handling in generator initialization.
The random number generator initialization could be more robust.
Consider this improvement:
def get_generator(
seed: int | list[int] | None = None,
) -> PaddleGenerator | None:
if seed is not None:
if isinstance(seed, list):
seed = mix_entropy(seed)
- if DEVICE == "cpu":
- generator = paddle.framework.core.default_cpu_generator()
- elif DEVICE == "gpu":
- generator = paddle.framework.core.default_cuda_generator(0)
- elif DEVICE.startswith("gpu:"):
- generator = paddle.framework.core.default_cuda_generator(
- int(DEVICE.split("gpu:")[1])
- )
- else:
- raise ValueError("DEVICE should be cpu or gpu or gpu:x")
+ try:
+ if DEVICE == "cpu":
+ generator = paddle.framework.core.default_cpu_generator()
+ elif DEVICE == "gpu":
+ generator = paddle.framework.core.default_cuda_generator(0)
+ elif DEVICE.startswith("gpu:"):
+ device_id = int(DEVICE.split("gpu:")[1])
+ if device_id < 0:
+ raise ValueError(f"Invalid GPU device ID: {device_id}")
+ generator = paddle.framework.core.default_cuda_generator(device_id)
+ else:
+ raise ValueError(f"Unsupported device: {DEVICE}. Expected 'cpu', 'gpu', or 'gpu:x'")
+ except Exception as e:
+ raise RuntimeError(f"Failed to initialize generator on device {DEVICE}: {str(e)}")
generator.manual_seed(seed)
return generator
else:
return None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def get_generator( | |
| seed: int | list[int] | None = None, | |
| ) -> PaddleGenerator | None: | |
| if seed is not None: | |
| if isinstance(seed, list): | |
| seed = mix_entropy(seed) | |
| if DEVICE == "cpu": | |
| generator = paddle.framework.core.default_cpu_generator() | |
| elif DEVICE == "gpu": | |
| generator = paddle.framework.core.default_cuda_generator(0) | |
| elif DEVICE.startswith("gpu:"): | |
| generator = paddle.framework.core.default_cuda_generator( | |
| int(DEVICE.split("gpu:")[1]) | |
| ) | |
| else: | |
| raise ValueError("DEVICE should be cpu or gpu or gpu:x") | |
| generator.manual_seed(seed) | |
| return generator | |
| else: | |
| return None | |
| def get_generator( | |
| seed: int | list[int] | None = None, | |
| ) -> PaddleGenerator | None: | |
| if seed is not None: | |
| if isinstance(seed, list): | |
| seed = mix_entropy(seed) | |
| try: | |
| if DEVICE == "cpu": | |
| generator = paddle.framework.core.default_cpu_generator() | |
| elif DEVICE == "gpu": | |
| generator = paddle.framework.core.default_cuda_generator(0) | |
| elif DEVICE.startswith("gpu:"): | |
| device_id = int(DEVICE.split("gpu:")[1]) | |
| if device_id < 0: | |
| raise ValueError(f"Invalid GPU device ID: {device_id}") | |
| generator = paddle.framework.core.default_cuda_generator(device_id) | |
| else: | |
| raise ValueError(f"Unsupported device: {DEVICE}. Expected 'cpu', 'gpu', or 'gpu:x'") | |
| except Exception as e: | |
| raise RuntimeError(f"Failed to initialize generator on device {DEVICE}: {str(e)}") | |
| generator.manual_seed(seed) | |
| return generator | |
| else: | |
| return None |
deepmd/pd/utils/nlist.py
Outdated
|
|
||
| # nb x nall x 3 | ||
| coord1 = coord.reshape([nb, -1, 3]) | ||
| nall = coord1.shape[1] |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove unused variable assignment.
The variable nall is assigned but never used.
Remove the unused variable:
- nall = coord1.shape[1]🧰 Tools
🪛 Ruff (0.8.0)
386-386: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/model/model/make_model.py
Outdated
| for kk in odef.keys(): | ||
| if kk not in model_ret.keys(): |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Optimize dictionary key membership checks.
Using .keys() for membership testing is inefficient.
Apply this optimization:
- for kk in odef.keys():
- if kk not in model_ret.keys():
+ for kk in odef:
+ if kk not in model_ret:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for kk in odef.keys(): | |
| if kk not in model_ret.keys(): | |
| for kk in odef: | |
| if kk not in model_ret: |
🧰 Tools
🪛 Ruff (0.8.0)
361-361: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
362-362: Use key not in dict instead of key not in dict.keys()
Remove .keys()
(SIM118)
deepmd/infer/__init__.py
Outdated
| "DeepPot", | ||
| "DeepPotential", |
There was a problem hiding this comment.
💡 Codebase verification
Add deprecation warning and documentation for DeepPotential
The code shows that DeepPotential is a factory function that forwards to DeepEval for compatibility reasons, while DeepPot is the main potential energy model class. The relationship should be documented:
- Add a deprecation warning in
DeepPotentialsince it exists only for compatibility - Add a note in docstring that
DeepPotis the recommended class to use - Update
__all__documentation to clarify the relationship
Example changes:
def DeepPotential(*args, **kwargs) -> "DeepEval":
"""Factory function that forwards to DeepEval (for compatibility).
.. deprecated:: 2.0
Use DeepPot instead. This function exists only for backwards compatibility.
Parameters
----------
*args
positional arguments
**kwargs
keyword arguments
Returns
-------
DeepEval
A DeepEval instance
"""
warnings.warn(
"DeepPotential is deprecated, use DeepPot instead",
DeprecationWarning,
stacklevel=2
)
return DeepEval(*args, **kwargs)🔗 Analysis chain
Consider documenting the relationship between DeepPot and DeepPotential
Both "DeepPot" and "DeepPotential" are now exposed in __all__. Since DeepPotential is a factory function that returns DeepEval for compatibility reasons, this might cause confusion for users. Consider:
- Adding a deprecation warning in
DeepPotentialif it's meant to be phased out - Adding documentation that explains the relationship between these classes
Let's check the usage patterns of these classes across the codebase:
Let me generate additional scripts to better understand the relationship between DeepPot and DeepPotential:
This will help us understand:
- The implementation of DeepPotential function
- Any docstrings or comments explaining its purpose
- The relationship with DeepPot class through docstrings
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check usage patterns of DeepPot vs DeepPotential
# to understand if both are actively used
# Search for DeepPot usage
echo "=== DeepPot usage ==="
rg "DeepPot[^e]" -A 2
# Search for DeepPotential usage
echo "=== DeepPotential usage ==="
rg "DeepPotential" -A 2
Length of output: 67928
Script:
#!/bin/bash
# Search for DeepPotential implementation and usage
echo "=== DeepPotential implementation ==="
ast-grep --pattern 'def DeepPotential($$$)'
echo -e "\n=== DeepPotential docstring ==="
rg -A 10 "def DeepPotential"
echo -e "\n=== DeepPot docstring ==="
rg -A 10 "class DeepPot"
Length of output: 7545
There was a problem hiding this comment.
Actionable comments posted: 55
🧹 Outside diff range and nitpick comments (101)
deepmd/dpmodel/fitting/__init__.py (1)
25-31: Maintain alphabetical ordering in__all__listConsider moving
make_base_fittingto maintain alphabetical ordering for better maintainability and consistency.__all__ = [ "DOSFittingNet", "DipoleFitting", "EnergyFittingNet", "InvarFitting", + "make_base_fitting", "PolarFitting", "PropertyFittingNet", - "make_base_fitting", ]deepmd/dpmodel/atomic_model/__init__.py (1)
51-56: LGTM! Consider grouping related models together.The reordering of items in
__all__improves readability. To further enhance organization, consider grouping related models together, for example:__all__ = [ "BaseAtomicModel", "DPAtomicModel", "DPDOSAtomicModel", "DPDipoleAtomicModel", "DPEnergyAtomicModel", "DPPolarAtomicModel", - "DPZBLLinearEnergyAtomicModel", - "LinearEnergyAtomicModel", + "LinearEnergyAtomicModel", + "DPZBLLinearEnergyAtomicModel", "PairTabAtomicModel", "make_base_atomic_model", ]This groups the linear energy models together while maintaining the overall structure.
deepmd/jax/env.py (1)
21-21: Consider documenting this public API addition.Since
jax_exportis now part of the public API, consider adding documentation that explains its purpose and usage patterns.Example docstring to add at the module level:
# SPDX-License-Identifier: LGPL-3.0-or-later +"""JAX environment configuration and exports. + +This module provides core JAX functionality including jax_export for compiling +and exporting JAX functions for deployment. +""" import osdeepmd/jax/descriptor/__init__.py (1)
28-35: Consider reorganizing the__all__listTo improve maintainability and readability, consider:
- Organizing the descriptors alphabetically
- Grouping related descriptors together (e.g., all Se-based descriptors)
Here's a suggested reorganization:
__all__ = [ - "DescrptDPA1", - "DescrptDPA2", - "DescrptHybrid", - "DescrptSeA", - "DescrptSeAttenV2", - "DescrptSeR", - "DescrptSeT", - "DescrptSeTTebd", + # DPA descriptors + "DescrptDPA2", + + # Se-based descriptors + "DescrptSeA", + "DescrptSeAttenV2", + "DescrptSeR", + "DescrptSeT", + "DescrptSeTTebd", ]deepmd/infer/model_devi.py (1)
207-214: Consider using f-strings consistentlyWhile the change to f-string on line 207 improves readability, the code could be further improved by using f-strings consistently throughout the header formatting.
Consider this improvement:
header = f"{header}\n{'step':10s}" for item in "vf": - header += "{:19s}{:19s}{:19s}".format( - f"max_devi_{item}", - f"min_devi_{item}", - f"avg_devi_{item}", - ) + header += f"{'max_devi_'+item:19s}{'min_devi_'+item:19s}{'avg_devi_'+item:19s}"deepmd/pt/cxx_op.py (1)
53-55: LGTM! Consider enhancing error messages with version-specific guidance.The transition to f-strings improves readability while maintaining the critical compatibility check information.
Consider adding more specific guidance in the error message for common scenarios. For example:
raise RuntimeError( f"This deepmd-kit package was compiled with " f"CXX11_ABI_FLAG={PT_CXX11_ABI_FLAG}, but PyTorch runtime was compiled " f"with CXX11_ABI_FLAG={pt_cxx11_abi_flag}. These two library ABIs are " f"incompatible and thus an error is raised when loading {module_name}. " - "You need to rebuild deepmd-kit against this PyTorch " - "runtime." + "You need to rebuild deepmd-kit against this PyTorch runtime. " + "For conda environments, ensure you install PyTorch from conda-forge. " + "For pip installations, try using `pip install --no-binary :all: deepmd-kit`." ) from edeepmd/env.py (1)
Line range hint
42-48: Enhance error message clarity for precision settings.Consider making the error message more explicit by listing all supported values in a more structured format.
raise RuntimeError( - f"Unsupported float precision option: {dp_float_prec}. Supported: high," - "low. Please set precision with environmental variable " - "DP_INTERFACE_PREC." + f"Unsupported float precision option: '{dp_float_prec}'\n" + "Supported values:\n" + "- 'high' or '' (default): double precision\n" + "- 'low': single precision\n" + "Please set precision using the DP_INTERFACE_PREC environment variable." )deepmd/dpmodel/fitting/dipole_fitting.py (1)
98-98: Add documentation for the new parameter.The newly added
dim_case_embdparameter is missing from the class docstring. Please document its purpose, type, and default value in the Parameters section.Add this to the class docstring's Parameters section:
Parameters ---------- + dim_case_embd : int, optional + The dimension of case embedding. Defaults to 0. rconddeepmd/pd/model/network/init.py (3)
56-56: Improve variable naming for better readabilityThe variable name
lis ambiguous and could be confused with the number '1'. Consider renaming it tolower_boundfor better clarity.- l = norm_cdf((a - mean) / std) + lower_bound = norm_cdf((a - mean) / std)🧰 Tools
🪛 Ruff (0.8.0)
56-56: Ambiguous variable name:
l(E741)
328-328: Add stacklevel to warning messagesThe warning messages for zero-element tensors should include an explicit stacklevel for better error tracing.
- warnings.warn("Initializing zero-element tensors is a no-op") + warnings.warn("Initializing zero-element tensors is a no-op", stacklevel=2)Also applies to: 378-378
🧰 Tools
🪛 Ruff (0.8.0)
328-328: No explicit
stacklevelkeyword argument found(B028)
1-456: Well-structured port from PyTorch to PaddleThe implementation successfully ports PyTorch's initialization module to Paddle while:
- Maintaining API compatibility
- Preserving mathematical correctness
- Following best practices for documentation
- Implementing proper error handling
This will help ensure consistent behavior between PyTorch and Paddle backends.
🧰 Tools
🪛 Ruff (0.8.0)
56-56: Ambiguous variable name:
l(E741)
328-328: No explicit
stacklevelkeyword argument found(B028)
378-378: No explicit
stacklevelkeyword argument found(B028)
deepmd/dpmodel/fitting/property_fitting.py (1)
81-81: Document the new parameter in the class docstring.The new parameter
dim_case_embdshould be documented in the class docstring to maintain API documentation completeness.Add the following to the Parameters section of the docstring:
type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + dim_case_embd: int, optional + Dimension for case embedding. Defaults to 0.deepmd/pd/utils/nlist.py (6)
18-25: Add type hints for input parametersConsider adding type hints for
coordandatypeparameters to improve code maintainability and IDE support.def extend_input_and_build_neighbor_list( - coord, - atype, + coord: paddle.Tensor, + atype: paddle.Tensor, rcut: float, sel: list[int], mixed_types: bool = False, box: Optional[paddle.Tensor] = None, ):
123-123: Consider usingtopkfor better performanceWhen selecting neighbors within the cutoff radius, using
paddle.topkinstead of full sorting could be more efficient since we only need the k nearest neighbors.- rr, nlist = paddle.sort(rr, axis=-1), paddle.argsort(rr, axis=-1) + values, indices = paddle.topk(rr, k=min(nsel, rr.shape[-1]), axis=-1, largest=False) + rr, nlist = values, indices
385-386: Remove unused variableThe variable
nallis not directly used and can be removed. The shape information can be accessed directly fromcoord1.- nall = coord1.shape[1]🧰 Tools
🪛 Ruff (0.8.0)
386-386: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
480-480: Clean up commented-out codeRemove the commented-out
.cpu()calls as they are no longer needed and may confuse future maintainers.dtype=env.GLOBAL_PD_FLOAT_PRECISION ) - # .cpu()Also applies to: 486-486, 492-492, 498-498
368-370: Add error handling for empty or invalid inputsConsider adding more comprehensive input validation to handle edge cases:
- Empty coordinate tensors
- Mismatched dimensions between coordinates and atom types
- Invalid cutoff radius values
if paddle.in_dynamic_mode(): assert len(rcuts) == len(nsels) + if not isinstance(rcuts, list) or not isinstance(nsels, list): + raise TypeError("rcuts and nsels must be lists") + if any(rc <= 0 for rc in rcuts): + raise ValueError("cutoff radius must be positive") if len(rcuts) == 0: return {}
397-400: Consider using batched operations for better performanceThe reshape and take_along_axis operations could be combined for better performance.
- coord2 = paddle.take_along_axis(coord1, axis=1, indices=index).reshape( - [nb, nloc, nsel, 3] - ) + coord2 = paddle.gather(coord1.reshape(nb, -1, 3), 1, index.reshape(nb, -1, 1).expand(-1, -1, 3)).reshape(nb, nloc, nsel, 3)deepmd/pd/model/descriptor/repformers.py (3)
441-442: Implement missing communication dictionary handlingThe code raises
NotImplementedErrorwhencomm_dictis provided, but there's a large commented-out implementation below. This suggests incomplete functionality.Would you like help implementing the communication dictionary handling based on the commented code? I can help create a GitHub issue to track this task.
548-552: Simplify conditional assignment using ternary operatorThe if-else block can be simplified using a ternary operator for better readability.
Apply this change:
- if callable(merged): - # only get data for once - sampled = merged() - else: - sampled = merged + # only get data for once if callable + sampled = merged() if callable(merged) else merged🧰 Tools
🪛 Ruff (0.8.0)
548-552: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
52-90: Add parameter validation in constructorThe constructor accepts many parameters but lacks validation for critical values like dimensions and layer counts.
Consider adding validation for:
- Positive values for dimensions (
g1_dim,g2_dim,axis_neuron)- Valid range for
nlayers(> 0)- Valid range for
rcutandrcut_smth(> 0 andrcut>rcut_smth)- Valid number of attention heads (
attn1_nhead,attn2_nhead> 0)Example validation:
def __init__(self, ...): if rcut <= rcut_smth: raise ValueError("rcut must be greater than rcut_smth") if g1_dim <= 0 or g2_dim <= 0: raise ValueError("dimensions must be positive") if nlayers <= 0: raise ValueError("nlayers must be positive")🧰 Tools
🪛 Ruff (0.8.0)
81-81: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/dpmodel/fitting/polarizability_fitting.py (1)
104-104: Documentation needs update for the new parameterThe addition of
dim_case_embdparameter is consistent with other fitting classes, but its purpose and usage should be documented in the class docstring.Add the following to the class docstring under Parameters section:
+ dim_case_embd : int, optional + The dimension of case embedding. Defaults to 0.Also applies to: 154-154
deepmd/dpmodel/fitting/invar_fitting.py (1)
126-126: Document the newdim_case_embdparameterThe new parameter
dim_case_embdis added but not documented in the class docstring. Please add its description to the Parameters section of the docstring.Add this to the docstring's Parameters section:
Parameters ---------- + dim_case_embd : int, optional + Dimension of the case embedding. Defaults to 0. bias_atomdeepmd/pt/model/descriptor/dpa2.py (5)
Line range hint
407-411: Simplify parameter sharing logicThe parameter sharing implementation has been simplified to two levels (0 and 1), which is more maintainable than the previous implementation. However, the code could benefit from more descriptive error messaging.
Consider enhancing the error message for unsupported shared levels:
- raise NotImplementedError + raise NotImplementedError( + f"Shared level {shared_level} is not supported. Only levels 0 (all parameters) " + "and 1 (type embedding only) are supported." + )
Line range hint
71-156: Consider adding validation for initialization parametersThe class initialization accepts many parameters but lacks input validation for critical values like
ntypes,rcut, etc.Add parameter validation to ensure robustness:
def __init__( self, ntypes: int, repinit: Union[RepinitArgs, dict], repformer: Union[RepformerArgs, dict], ... ) -> None: + if ntypes <= 0: + raise ValueError("ntypes must be positive") + if not isinstance(repinit, (dict, RepinitArgs)): + raise TypeError("repinit must be a dict or RepinitArgs") + if not isinstance(repformer, (dict, RepformerArgs)): + raise TypeError("repformer must be a dict or RepformerArgs")
Line range hint
158-200: Add docstring forshare_paramsmethodThe
share_paramsmethod lacks documentation explaining the purpose and behavior of different sharing levels.Add a comprehensive docstring:
def share_params(self, base_class, shared_level, resume=False) -> None: + """Share parameters between descriptors for multi-task training. + + Parameters + ---------- + base_class : DescrptDPA2 + The base descriptor class to share parameters with + shared_level : int + The level of parameter sharing: + - 0: Share all parameters (type_embedding, repinit, repformers) + - 1: Share only type_embedding parameters + resume : bool, optional + If True, resume from checkpoint without recalculating statistics + + Raises + ------ + NotImplementedError + If shared_level is not 0 or 1 + AssertionError + If base_class is not of the same type + """
Line range hint
202-264: Enhance error handling inchange_type_mapThe
change_type_mapmethod could benefit from more robust error handling and validation.Add validation for input parameters and handle edge cases:
def change_type_map( self, type_map: list[str], model_with_new_type_stat=None ) -> None: + if not isinstance(type_map, list) or not all(isinstance(t, str) for t in type_map): + raise TypeError("type_map must be a list of strings") + if len(type_map) == 0: + raise ValueError("type_map cannot be empty") + if len(set(type_map)) != len(type_map): + raise ValueError("type_map contains duplicate types") assert self.type_map is not None, "'type_map' must be defined when performing type changing!"
Line range hint
266-300: Consider optimizing memory usage in forward passThe
forwardmethod creates several intermediate tensors that could be reused or optimized.Consider implementing memory optimization techniques:
- Use in-place operations where possible
- Reuse intermediate tensors
- Consider implementing a memory pool for frequently allocated tensors of similar sizes
deepmd/pd/utils/env.py (1)
79-151: LGTM: Well-implemented primitive API configurationThe implementation is thorough with proper operator blacklisting and logging. Consider adding a brief explanation of why specific operators are blacklisted for better maintainability.
Add a comment explaining the criteria for including operators in the blacklist.
deepmd/pd/model/descriptor/se_t_tebd.py (3)
451-455: Remove unused variable assignmentThe variable
nallis assigned but never used in this code block.Apply this fix:
# cast the input to internal precsion extended_coord = extended_coord.to(dtype=self.prec) del mapping nframes, nloc, nnei = nlist.shape -nall = extended_coord.reshape([nframes, -1]).shape[1] // 3 g1_ext = self.type_embedding(extended_atype)🧰 Tools
🪛 Ruff (0.8.0)
454-454: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
🪛 GitHub Check: CodeQL
[notice] 451-451: Unused local variable
Variable nall is not used.
823-828: Consider using einsum for better readability and potential performanceThe current tensor operation could be more clearly expressed using einsum notation.
Consider this alternative implementation:
-env_ij = ( - # ij1m x i1km -> ijkm -> ijk - rr_i.unsqueeze(2) * rr_j.unsqueeze(1) -).sum(-1) +env_ij = paddle.einsum("ijm,ikm->ijk", rr_i, rr_j)This change makes the intention clearer and might be optimized better by the framework.
909-914: Consider using einsum for better readabilitySimilar to the previous suggestion, this tensor operation could be more clearly expressed using einsum notation.
Consider this alternative implementation:
-res_ij = ( - # ijk1 x ijkm -> ijkm -> im - env_ij.unsqueeze(-1) * gg -).sum([1, 2]) +res_ij = paddle.einsum("ijk,ijkm->im", env_ij, gg)deepmd/common.py (1)
Line range hint
47-60: Strong type safety foundation laid for precision and activation functions.The introduction of Literal types and type-safe sets provides several benefits:
- Early error detection for invalid precision/activation values
- Better IDE support with autocomplete
- Improved static type checking capabilities
Consider extending this pattern to other configuration options in the codebase for consistent type safety.
deepmd/pd/utils/utils.py (4)
35-58: Optimize activation function selection performanceThe current implementation calls
.lower()onself.activationin everyforwardcall. Consider storing the lowercase version during initialization to avoid repeated string transformations.def __init__(self, activation: str | None): super().__init__() - self.activation: str = activation if activation is not None else "linear" + self.activation: str = (activation if activation is not None else "linear").lower() def forward(self, x: paddle.Tensor) -> paddle.Tensor: """Returns the tensor after applying activation function corresponding to `activation`.""" - if self.activation.lower() == "relu": + if self.activation == "relu": return F.relu(x) - elif self.activation.lower() == "gelu" or self.activation.lower() == "gelu_tf": + elif self.activation == "gelu" or self.activation == "gelu_tf": return F.gelu(x, approximate=True)🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 55-55: Statement has no effect
This statement has no effect.
71-73: Remove redundant assertionThe assertion
assert xx is not Noneis redundant as it's already guaranteed by the previous if-return statement.if xx is None: return None -assert xx is not None
138-164: Enhance documentation for hash mixing functionsWhile the implementation correctly follows numpy's approach, it would be beneficial to add docstrings explaining:
- The purpose and importance of hash mixing
- The mathematical properties of the constants used
- The relationship between
hashmix,mix, andmix_entropy
188-198: Add error handling documentation for nvprof_contextConsider adding a docstring that explains:
- The purpose of NVTX profiling
- What exceptions might be raised by
nvprof_nvtx_pushandnvprof_nvtx_pop- How errors are handled within the context manager
deepmd/pd/model/model/make_model.py (2)
361-362: Optimize dictionary key lookup operations.Replace
.keys()with direct dictionary membership testing for better performance:-if kk not in odef.keys(): +if kk not in odef:🧰 Tools
🪛 Ruff (0.8.0)
361-361: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
362-362: Use
key not in dictinstead ofkey not in dict.keys()Remove
.keys()(SIM118)
440-476: Consider optimizing neighbor list formatting performance.The current implementation of neighbor list formatting involves multiple tensor operations that could be optimized:
- The padding operation uses
paddle.concatwhich creates a new tensor- Multiple tensor operations for sorting and filtering could potentially be combined
Consider pre-allocating the padded tensor and using in-place operations where possible:
- nlist = paddle.concat( - [ - nlist, - -1 - * paddle.ones( - [n_nf, n_nloc, nnei - n_nnei], - dtype=nlist.dtype, - ).to(nlist.place), - ], - axis=-1, - ) + padded_nlist = paddle.full([n_nf, n_nloc, nnei], -1, + dtype=nlist.dtype).to(nlist.place) + padded_nlist[..., :n_nnei] = nlist + nlist = padded_nlistdeepmd/dpmodel/atomic_model/dp_atomic_model.py (1)
68-73: Consider adding error handling for invalid case indices.The method should validate the
case_idxparameter before passing it to the fitting object to prevent potential runtime errors.def set_case_embd(self, case_idx: int): """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ + if not isinstance(case_idx, int) or case_idx < 0: + raise ValueError(f"case_idx must be a non-negative integer, got {case_idx}") self.fitting.set_case_embd(case_idx)deepmd/pt/model/model/make_model.py (1)
517-519: Add docstring for clarityWhile the implementation is correct, consider adding a docstring to document the method's purpose and parameters.
def set_case_embd(self, case_idx: int): + """Set the case embedding index in the atomic model. + + Parameters + ---------- + case_idx : int + The case embedding index to be set in the atomic model. + """ self.atomic_model.set_case_embd(case_idx)deepmd/infer/deep_eval.py (1)
506-506: LGTM: Error message improvementsThe updates to error messages using f-strings improve readability and maintainability. The messages clearly communicate the expected dimensions for both frame parameters and atomic parameters.
However, there's a minor inconsistency in the error message on line 518:
- For
aparam, the message says "frame param" but should say "atomic param" to match the parameter type.- f"got wrong size of frame param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}" + f"got wrong size of atomic param, should be either {nframes} x {natoms} x {fdim} or {natoms} x {fdim} or {fdim}"Also applies to: 518-518
deepmd/pd/model/descriptor/dpa2.py (4)
499-499: Rename unused loop variableiito_The loop variable
iiis not used within the loop body. Renaming it to_clarifies that it is intentionally unused.Apply this diff:
-for ii, descrpt in enumerate(descrpt_list): +for _, descrpt in enumerate(descrpt_list):🧰 Tools
🪛 Ruff (0.8.0)
499-499: Loop control variable
iinot used within loop bodyRename unused
iito_ii(B007)
646-646: Correct the misspelling ofg1_shape_tranformThe attribute
g1_shape_tranformappears to be misspelled. It should beg1_shape_transformto maintain consistency and readability.Apply this diff:
-obj.g1_shape_tranform = MLPLayer.deserialize(g1_shape_tranform) +obj.g1_shape_transform = MLPLayer.deserialize(g1_shape_transform)Ensure to update all occurrences of
g1_shape_tranformin the codebase.
663-663: Remove unused variableenv_matassignmentsThe variable
env_matis assigned but never used in the deserialization methods. This is unnecessary and can be removed.Apply this diff:
-env_mat = repinit_variable.pop("env_mat")-env_mat = repinit_three_body_variable.pop("env_mat")-env_mat = repformers_variable.pop("env_mat")Also applies to: 678-678, 684-684
🧰 Tools
381-383: Improve documentation inshare_paramsmethodThe docstring for the
share_paramsmethod lacks clarity and proper formatting. Enhancing it can improve readability and maintainability.Consider reformatting the docstring:
""" Share the parameters of `self` with `base_class` at the specified `shared_level` during multitask training. If not starting from a checkpoint (`resume=False`), some separated parameters (e.g., mean and stddev) will be recalculated across different instances. """deepmd/pd/train/wrapper.py (1)
30-34: Fix incomplete and incorrect docstring.The docstring references a non-existent
configparameter and is missing documentation for other parameters.Update the docstring to:
- """Construct a DeePMD model wrapper. - - Args: - - config: The Dict-like configuration with training options. - """ + """Construct a DeePMD model wrapper. + + Args: + model: Single model layer or dictionary of model layers for multi-task + loss: Optional loss layer or dictionary of loss layers + model_params: Optional model parameters dictionary + shared_links: Optional parameter sharing configuration + """deepmd/pd/utils/decomp.py (2)
70-92: Consider improving variable names and validation.The function is well-implemented but could benefit from more descriptive variable names and additional validation.
Consider these improvements:
-def sec(length: int, size: int) -> list[int]: +def chunk_sizes(total_length: int, chunk_size: int) -> list[int]: """Auxiliary function for decomposed functions. If length is not divisible by size, the last chunk will be smaller. Parameters ---------- - length : int + total_length : int Length to be chunked. - size : int + chunk_size : int Chunk size. Returns ------- list[int] Chunked output list. """ - assert length > 0 - assert size > 0 - if length % size == 0: - return [size] * (length // size) - return [size] * (length // size) + [length % size] + if not isinstance(total_length, int) or not isinstance(chunk_size, int): + raise TypeError("Both arguments must be integers") + if total_length <= 0: + raise ValueError(f"total_length must be positive, got {total_length}") + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + + num_full_chunks = total_length // chunk_size + remainder = total_length % chunk_size + + return [chunk_size] * num_full_chunks + ([remainder] if remainder else [])
130-134: Consider optimizing dynamic mode handling.The function could be improved in terms of validation and efficiency.
Consider these improvements:
def numel(x: paddle.Tensor) -> int: + if not isinstance(x, paddle.Tensor): + raise TypeError(f"Expected paddle.Tensor, got {type(x)}") + if paddle.in_dynamic_mode(): - return np.prod(x.shape) + return paddle.prod(paddle.to_tensor(x.shape)).item() return paddle.numel(x)deepmd/pd/model/task/invar_fitting.py (2)
128-130: Simplify the remove_vaccum_contribution conditionThe condition for
remove_vaccum_contributionis complex and could be simplified for better readability.Consider refactoring to:
- remove_vaccum_contribution=None - if atom_ener is None or len([x for x in atom_ener if x is not None]) == 0 - else [x is not None for x in atom_ener], + remove_vaccum_contribution=self._get_remove_vaccum_contribution(atom_ener),And add a helper method:
def _get_remove_vaccum_contribution(self, atom_ener): if not atom_ener: return None has_valid_energies = any(x is not None for x in atom_ener) return [x is not None for x in atom_ener] if has_valid_energies else None
176-176: Fix typo in docstringThere's a typo in the docstring: "alculate" should be "calculate".
deepmd/pd/model/model/__init__.py (3)
43-64: Consider optimizing type_map copying and adding parameter validation.The function makes redundant deep copies of
type_mapand lacks validation for required parameters.Consider these improvements:
def _get_standard_model_components(model_params, ntypes): + # Validate required parameters + if "type_map" not in model_params: + raise ValueError("type_map is required in model_params") + if "descriptor" not in model_params: + raise ValueError("descriptor is required in model_params") + + # Copy type_map once and reuse + type_map_copy = copy.deepcopy(model_params["type_map"]) + # descriptor model_params["descriptor"]["ntypes"] = ntypes - model_params["descriptor"]["type_map"] = copy.deepcopy(model_params["type_map"]) + model_params["descriptor"]["type_map"] = type_map_copy descriptor = BaseDescriptor(**model_params["descriptor"]) # fitting fitting_net = model_params.get("fitting_net", {}) fitting_net["type"] = fitting_net.get("type", "ener") fitting_net["ntypes"] = descriptor.get_ntypes() - fitting_net["type_map"] = copy.deepcopy(model_params["type_map"]) + fitting_net["type_map"] = type_map_copy
75-95: Consider simplifying the bias conversion logic.The nested loops and conditionals could be simplified for better readability and maintainability.
Consider this more concise approach:
def _convert_preset_out_bias_to_array(preset_out_bias, type_map): if preset_out_bias is not None: for kk in preset_out_bias: if len(preset_out_bias[kk]) != len(type_map): raise ValueError( - "length of the preset_out_bias should be the same as the type_map" + f"Length mismatch: preset_out_bias['{kk}'] has length {len(preset_out_bias[kk])}, " + f"but type_map has length {len(type_map)}" ) - for jj in range(len(preset_out_bias[kk])): - if preset_out_bias[kk][jj] is not None: - if isinstance(preset_out_bias[kk][jj], list): - bb = preset_out_bias[kk][jj] - elif _can_be_converted_to_float(preset_out_bias[kk][jj]): - bb = [float(preset_out_bias[kk][jj])] - else: - raise ValueError( - f"unsupported type/value of the {jj}th element of " - f"preset_out_bias['{kk}'] " - f"{type(preset_out_bias[kk][jj])}" - ) - preset_out_bias[kk][jj] = np.array(bb) + preset_out_bias[kk] = [ + np.array(v if isinstance(v, list) else [float(v)]) + if v is not None and (isinstance(v, list) or _can_be_converted_to_float(v)) + else v + for v in preset_out_bias[kk] + ] return preset_out_bias
129-135: Enhance error handling in the model factory.The error handling in the model factory could be more informative.
Consider this improvement:
def get_model(model_params): + if not isinstance(model_params, dict): + raise TypeError("model_params must be a dictionary") + model_type = model_params.get("type", "standard") if model_type == "standard": return get_standard_model(model_params) else: - return BaseModel.get_class_by_type(model_type).get_model(model_params) + try: + model_class = BaseModel.get_class_by_type(model_type) + return model_class.get_model(model_params) + except (KeyError, AttributeError) as e: + raise ValueError(f"Invalid model type: {model_type}") from edeepmd/pd/utils/dataloader.py (3)
86-86: Simplify iteration overh5py.FileobjectSince
h5py.Fileobjects are directly iterable over their keys, you can simplify the loop by iterating overfileinstead offile.keys(). This improves readability.- systems = [os.path.join(systems, item) for item in file.keys()] + systems = [os.path.join(systems, item) for item in file]🧰 Tools
🪛 Ruff (0.8.0)
86-86: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
285-290: Combine nestedifstatements for clarityThe nested
ifstatements can be combined into a singleifstatement using logicalandoperators. This enhances readability and reduces unnecessary indentation.- if self._queue.qsize() < min(2, max(1, self._queue.maxsize // 2)): - if time.time() - self.start_time > 5 * 60: - if ( - self.warning_time is None - or time.time() - self.warning_time > 15 * 60 - ): + if ( + self._queue.qsize() < min(2, max(1, self._queue.maxsize // 2)) + and time.time() - self.start_time > 5 * 60 + and ( + self.warning_time is None + or time.time() - self.warning_time > 15 * 60 + ) + ):🧰 Tools
🪛 Ruff (0.8.0)
285-290: Use a single
ifstatement instead of nestedifstatements(SIM102)
286-290: Use a single
ifstatement instead of nestedifstatementsCombine
ifstatements usingand(SIM102)
310-310: Simplify dictionary iteration by omitting.keys()When iterating over a dictionary, you can iterate directly over the dictionary without calling
.keys(). This is more idiomatic in Python and improves readability.- for key in example.keys(): + for key in example:🧰 Tools
🪛 Ruff (0.8.0)
310-310: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
deepmd/pd/model/descriptor/se_atten.py (3)
375-379: Simplify assignment using a ternary operatorConsider using a ternary operator for conciseness and readability.
Apply this diff:
if callable(merged): # only get data for once - sampled = merged() else: - sampled = merged + sampled = merged() if callable(merged) else merged🧰 Tools
🪛 Ruff (0.8.0)
375-379: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
477-477: Remove unused variable 'nall'The local variable
nallis assigned but never used, which can be removed to clean up the code.Apply this diff:
- nall = extended_coord.reshape([nb, -1, 3]).shape[1]🧰 Tools
🪛 Ruff (0.8.0)
477-477: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
816-818: Ensure consistent seed usage for reproducibilityIn the initialization of
GatedAttentionLayer, confirm that seeds are consistently used to maintain reproducibility across different runs.Apply this diff to pass a child seed:
self.attention_layer = GatedAttentionLayer( nnei, embed_dim, hidden_dim, dotr=dotr, do_mask=do_mask, scaling_factor=scaling_factor, normalize=normalize, temperature=temperature, smooth=smooth, precision=precision, - seed=child_seed(seed, 0), + seed=child_seed(self.seed, 0), )deepmd/pd/model/task/fitting.py (4)
328-328: Simplify loop by iterating directly over the dictionaryWhen iterating over a dictionary, you can iterate directly over its keys without calling
.keys(). This improves readability and is more idiomatic in Python.Apply this diff to simplify the loop:
- for kk in variables.keys(): + for kk in variables:🧰 Tools
🪛 Ruff (0.8.0)
328-328: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
435-439: Address the TODO: Compute input for vacuum contributionThere's a TODO comment indicating that the computation for the vacuum contribution input when
remove_vaccum_contributionis set is pending. Implementing this is important for accurate calculations when vacuum contributions are involved.Do you need assistance in implementing this functionality, or would you like me to open a new GitHub issue to track this task?
513-548: Optimize repeated condition checks in the forward passIn the
_forward_commonmethod, there are repeated condition checks forxx_zeros is not Noneinside loops. This can be optimized by checking once before the loops to improve performance.Consider restructuring the code to minimize repeated condition checks. For example, you can move the check outside the loop or store the result in a variable to avoid redundant evaluations.
363-369: Type hint forcase_idxinset_case_embdmethodIn the
set_case_embdmethod, adding a type hint forcase_idxcan improve code clarity and static analysis.Add the type hint for
case_idx:- def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None:deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
266-266: Typo in docstring: Correct 'mapps' to 'maps'In the docstring for
forward_atomic, there is a typo:'mapps'should be'maps'.Apply this diff to fix the typo:
mapping - mapps the extended indices to local indices + maps the extended indices to local indicesdeepmd/pd/model/descriptor/dpa1.py (3)
534-534: Remove unused variableenv_matThe variable
env_matis assigned but never used. Removing this assignment will clean up the code.Apply this diff to fix the issue:
env_mat = data.pop("env_mat")🧰 Tools
🪛 Ruff (0.8.0)
534-534: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
633-633: Remove unused variablenallThe variable
nallis assigned but not used anywhere in theforwardmethod. Removing it will eliminate unnecessary code.Apply this diff to fix the issue:
nframes, nloc, nnei = nlist.shape - nall = extended_coord.reshape([nframes, -1]).shape[1] // 3 g1_ext = self.type_embedding(extended_atype)🧰 Tools
🪛 Ruff (0.8.0)
633-633: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
604-604: Typo in docstring: Correct 'aotm' to 'atom'In the docstring for
forward, there is a typo:'aotm'should be'atom'.Apply this diff to fix the typo:
extended_atype - The extended aotm types. shape: nf x nall + The extended atom types. shape: nf x nalldeepmd/pt/model/task/fitting.py (2)
361-369: Add bounds checking for 'case_idx' in 'set_case_embd' methodIn the
set_case_embdmethod, there's no validation to ensure thatcase_idxis within the valid range (0 ≤case_idx<dim_case_embd). An invalidcase_idxcould cause anIndexError. Adding bounds checking will enhance the robustness of the method.Apply this diff to add validation:
def set_case_embd(self, case_idx: int): """ Set the case embedding of this fitting net by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ + if not (0 <= case_idx < self.dim_case_embd): + raise ValueError(f"'case_idx' must be between 0 and {self.dim_case_embd - 1}") self.case_embd = torch.eye(self.dim_case_embd, dtype=self.prec, device=device)[ case_idx ]
500-512: Optimize 'case_embd' replication using broadcasting instead of 'torch.tile'In the
_forward_commonmethod,torch.tileis used to replicatecase_embdacross frames and atoms, which may impact performance and memory usage. You can achieve the same result more efficiently by reshapingcase_embdand leveraging broadcasting.Apply this diff to use broadcasting:
if self.dim_case_embd > 0: assert self.case_embd is not None - case_embd = torch.tile(self.case_embd.reshape([1, 1, -1]), [nf, nloc, 1]) + case_embd = self.case_embd.view(1, 1, -1) xx = torch.cat( [xx, case_embd], dim=-1, ) if xx_zeros is not None: - xx_zeros = torch.tile(xx_zeros, [nf, nloc, 1]) + xx_zeros = xx_zeros.expand(nf, nloc, -1) xx_zeros = torch.cat( [xx_zeros, case_embd], dim=-1, )deepmd/dpmodel/fitting/general_fitting.py (2)
Line range hint
108-131: Add input validation fordim_case_embdparameterConsider adding a check to ensure that
dim_case_embdis a non-negative integer. This will prevent potential issues if an invalid value is passed.Apply this diff to add input validation:
def __init__( ... dim_case_embd: int = 0, bias_atom_e: Optional[np.ndarray] = None, ... ): + if dim_case_embd < 0: + raise ValueError(f"'dim_case_embd' must be non-negative, got {dim_case_embd}") self.dim_case_embd = dim_case_embd ...
Line range hint
308-325: Handle versioning indeserializemethodSince the version is incremented to 3 in
serialize, consider updating thedeserializemethod to handle different versions and maintain backward compatibility.deepmd/pd/model/network/network.py (4)
206-206: Fix typo in docstring: 'Caulate' should be 'Calculate'There's a typo in the docstring of the
forwardmethod inTypeEmbedNetConsistent. Correcting it will enhance code readability.Apply this diff to fix the typo:
- """Caulate type embedding network. + """Calculate type embedding network.
76-80: Complete the Returns section in the docstringThe
Returnssection in the docstring of theforwardmethod inTypeEmbedNetis incomplete. Adding a clear description will improve documentation clarity.Consider updating the Returns section:
Returns ------- -type_embedding: +type_embedding: paddle.Tensor + The embedding tensor corresponding to the input atom types.
205-230: Refactorforwardmethod to accept input data instead of deviceThe
forwardmethod inTypeEmbedNetConsistentacceptsdeviceas an argument and does not process input data directly, which is unconventional for neural network modules. Typically,forwardmethods accept input tensors and perform computations on them.Consider modifying the
forwardmethod to accept input data (e.g., atom types) and handle device placement internally to align with standard conventions.
30-34: Remove unused functionTensorThe function
Tensoris defined but not used anywhere in the code. Removing unused code improves maintainability and readability.Consider removing the unused function:
-def Tensor(*shape): - return paddle.empty(shape, dtype=env.GLOBAL_PD_FLOAT_PRECISION).to( - device=env.DEVICE - )deepmd/pd/model/descriptor/se_a.py (4)
361-361: Replace unused variableenv_matwith_to indicate intentional discardThe variable
env_matis assigned but never used. Assigning it to_indicates that the discard is intentional and prevents linters from flagging it.- env_mat = data.pop("env_mat") + _ = data.pop("env_mat")🧰 Tools
🪛 Ruff (0.8.0)
361-361: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
655-655: Remove unnecessary assignment totiThe variable
tiis assigned a value but is immediately reassigned without being used. Removing the initial assignment avoids confusion.- ti = -1🧰 Tools
🪛 GitHub Check: CodeQL
[warning] 655-655: Variable defined multiple times
This assignment to 'ti' is unnecessary as it is redefined before this value is used.
717-717: Rename unused loop variablescompress_data_iiandcompress_info_iito_The loop variables
compress_data_iiandcompress_info_iiare not used within the loop body. Renaming them to_clarifies that they are intentionally unused.- for embedding_idx, (ll, compress_data_ii, compress_info_ii) in enumerate(zip(self.filter_layers.networks, self.compress_data, self.compress_info)): + for embedding_idx, (ll, _, _) in enumerate(zip(self.filter_layers.networks, self.compress_data, self.compress_info)):🧰 Tools
🪛 Ruff (0.8.0)
717-717: Loop control variable
compress_data_iinot used within loop bodyRename unused
compress_data_iito_compress_data_ii(B007)
717-717: Loop control variable
compress_info_iinot used within loop bodyRename unused
compress_info_iito_compress_info_ii(B007)
47-52: Consider adding a comment or usingcontextlib.suppressfor the emptyexceptblockThe
exceptblock only containspasswithout any explanatory comment, which may lead to confusion. Adding a comment or usingcontextlib.suppressmakes the intent clearer.Option 1: Add an explanatory comment:
try: from typing import ( Final, ) except ImportError: + # `Final` is not available in this version of Python; skipping import. passOption 2: Use
contextlib.suppress:+from contextlib import suppress ... with suppress(ImportError): from typing import ( Final, )🧰 Tools
🪛 Ruff (0.8.0)
47-52: Use
contextlib.suppress(ImportError)instead oftry-except-passReplace with
contextlib.suppress(ImportError)(SIM105)
🪛 GitHub Check: CodeQL
[notice] 50-50: Empty except
'except' clause does nothing but pass and there is no explanatory comment.deepmd/pd/loss/ener.py (1)
145-347: Consider refactoring theforwardmethod to reduce code duplicationThe
forwardmethod is quite lengthy and contains repeated code patterns for computing losses for different components (energy, force, virial, etc.). To improve maintainability and readability, consider refactoring the method to extract common operations into helper methods. This will reduce code duplication and make the code easier to understand and maintain.deepmd/pd/utils/region.py (4)
27-27: Remove unused variableein the exception blockThe variable
eis assigned but never used in theexceptclause. Removing it will clean up the code.Apply this diff to fix the issue:
- except Exception as e: + except Exception:🧰 Tools
🪛 Ruff (0.8.0)
27-27: Local variable
eis assigned to but never usedRemove assignment to unused variable
e(F841)
78-78: Add type annotations to function parameters for consistencyThe function
b_to_face_distancelacks type annotations for its parameters and return type, unlike other functions in the file. Adding type annotations improves code readability and consistency.Apply this diff to add type annotations:
-def b_to_face_distance(cell): +def b_to_face_distance(cell: paddle.Tensor) -> paddle.Tensor:
89-91: Remove commented-out code to clean up the fileLines 89-90 contain commented-out code that doesn't appear to be needed. Removing it will improve code readability.
Apply this diff to remove the commented-out code:
-# b_to_face_distance = paddle.vmap( -# _to_face_distance, in_dims=(0), out_dims=(0))
97-108: Document all function parameters in the docstringIn the
normalize_coordfunction, the parametercellis not documented in the docstring. Adding documentation for all parameters improves code clarity.Apply this diff to add the missing parameter documentation:
""" Apply PBC according to the atomic coordinates. Parameters ---------- coord : paddle.Tensor original coordinates of shape [*, na, 3]. + cell : paddle.Tensor + simulation cell tensor of shape [*, 3, 3]. Returns ------- wrapped_coord: paddle.Tensor wrapped coordinates of shape [*, na, 3]. """deepmd/pd/entrypoints/main.py (5)
118-118: Remove unused variablefin the context managerThe variable
fis assigned but never used. Since we're only creating an empty file, we can omitffrom thewithstatement.Apply this diff to fix the issue:
- with h5py.File(stat_file_path_single, "w") as f: + with h5py.File(stat_file_path_single, "w"):🧰 Tools
🪛 Ruff (0.8.0)
118-118: Local variable
fis assigned to but never usedRemove assignment to unused variable
f(F841)
500-503: Simplify theif-elseblock using a ternary operatorYou can replace the
if-elseblock with a single line using a ternary operator to make the code more concise.Apply this diff to simplify the code:
- if not isinstance(args, argparse.Namespace): - FLAGS = parse_args(args=args) - else: - FLAGS = args + FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args🧰 Tools
🪛 Ruff (0.8.0)
500-503: Use ternary operator
FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else argsinstead ofif-else-blockReplace
if-else-block withFLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args(SIM108)
Line range hint
645-647: Use a context manager to ensure the file is properly closedCurrently, the file
fout1may not be properly closed if an exception occurs before reachingfout1.close(). Using a context manager (withstatement) ensures that the file is closed correctly.Refactor the code to use a
withstatement:if SAMPLER_RECORD: record_file = f"Sample_rank_{self.rank}.txt" - fout1 = open(record_file, mode="w", buffering=1) + with open(record_file, mode="w", buffering=1) as fout1: + # Rest of the code that uses fout1 + ...You will need to adjust the code to ensure that all usages of
fout1are within the context manager.🧰 Tools
🪛 Ruff (0.8.0)
118-118: Local variable
fis assigned to but never usedRemove assignment to unused variable
f(F841)
500-503: Use ternary operator
FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else argsinstead ofif-else-blockReplace
if-else-block withFLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args(SIM108)
Line range hint
636-640: Use a context manager to ensure the file is properly closedCurrently, the file
foutmay not be properly closed if an exception occurs before reachingfout.close(). Using a context manager (withstatement) ensures that the file is closed correctly.Refactor the code to use a
withstatement:- fout = ( - open( - self.disp_file, - mode="w" if not self.restart_training else "a", - buffering=1, - ) - if self.rank == 0 - else None - ) # line buffered + if self.rank == 0: + with open( + self.disp_file, + mode="w" if not self.restart_training else "a", + buffering=1, + ) as fout: + # Rest of the code that uses fout + ...You will need to adjust the code to ensure that all usages of
foutare within the context manager.🧰 Tools
🪛 Ruff (0.8.0)
118-118: Local variable
fis assigned to but never usedRemove assignment to unused variable
f(F841)
500-503: Use ternary operator
FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else argsinstead ofif-else-blockReplace
if-else-block withFLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args(SIM108)
Line range hint
683-686: Simplify theif-elseblock using a ternary operatorYou can replace the
if-elseblock with a single line using a ternary operator to make the code more concise.Apply this diff to simplify the code:
if _step_id < self.warmup_steps: pref_lr = _lr.start_lr else: pref_lr = cur_lr + # Simplify the above using a ternary operator + pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr🧰 Tools
🪛 Ruff (0.8.0)
118-118: Local variable
fis assigned to but never usedRemove assignment to unused variable
f(F841)
500-503: Use ternary operator
FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else argsinstead ofif-else-blockReplace
if-else-block withFLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args(SIM108)
deepmd/pd/train/training.py (8)
290-290: Simplifydict.getmethod usageYou can simplify
config.get("learning_rate_dict", None)toconfig.get("learning_rate_dict"), as the default value isNone.Apply this diff to simplify the code:
- if config.get("learning_rate_dict", None) is not None: + if config.get("learning_rate_dict") is not None:🧰 Tools
🪛 Ruff (0.8.0)
290-290: Use
config.get("learning_rate_dict")instead ofconfig.get("learning_rate_dict", None)Replace
config.get("learning_rate_dict", None)withconfig.get("learning_rate_dict")(SIM910)
392-392: Simplifydict.getmethod usageYou can simplify
config.get("learning_rate_dict", None)toconfig.get("learning_rate_dict"), as the default value isNone.Apply this diff to simplify the code:
if self.multi_task and config.get("learning_rate_dict", None) is not None: + if self.multi_task and config.get("learning_rate_dict") is not None:🧰 Tools
🪛 Ruff (0.8.0)
392-392: Use
config.get("learning_rate_dict")instead ofconfig.get("learning_rate_dict", None)Replace
config.get("learning_rate_dict", None)withconfig.get("learning_rate_dict")(SIM910)
493-493: Simplify dictionary iteration by removing unnecessary.keys()When iterating over a dictionary, you can omit
.keys()since iterating over the dictionary itself yields its keys.Apply this diff to simplify the code:
target_keys = [ i - for i in _random_state_dict.keys() + for i in _random_state_dict if i != "_extra_state" and f".{_model_key}." in i ]🧰 Tools
🪛 Ruff (0.8.0)
493-493: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
683-686: Simplify theif-elseblock using a ternary operatorYou can replace the
if-elseblock with a single line using a ternary operator to make the code more concise.Apply this diff to simplify the code:
if _step_id < self.warmup_steps: pref_lr = _lr.start_lr else: pref_lr = cur_lr + # Simplify the above using a ternary operator + pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr🧰 Tools
🪛 Ruff (0.8.0)
683-686: Use ternary operator
pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lrinstead ofif-else-blockReplace
if-else-block withpref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr(SIM108)
739-739: Rename unused loop control variableThe loop control variable
iiis not used within the loop body. To indicate that it's intentionally unused, consider renaming it to_.Apply this diff to rename the variable:
for ii in range(valid_numb_batch): + for _ in range(valid_numb_batch):🧰 Tools
🪛 Ruff (0.8.0)
739-739: Loop control variable
iinot used within loop bodyRename unused
iito_ii(B007)
863-863: Remove unused local variablemoduleThe variable
moduleis assigned but never used. Removing it will clean up the code.Apply this diff to remove the unused variable:
- module = ( - self.wrapper.module - if dist.is_available() and dist.is_initialized() - else self.wrapper - )🧰 Tools
🪛 Ruff (0.8.0)
863-863: Local variable
moduleis assigned to but never usedRemove assignment to unused variable
module(F841)
1048-1048: Simplify dictionary iteration by removing unnecessary.keys()When iterating over a dictionary, you can omit
.keys()since iterating over the dictionary itself yields its keys.Apply this diff to simplify the code:
for key in batch_data.keys(): + for key in batch_data:🧰 Tools
🪛 Ruff (0.8.0)
1048-1048: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
Line range hint
1278-1278: Remove unused local variablenitemThe variable
nitemis assigned but never used. Removing it will clean up the code.Apply this diff to remove the unused variable:
def list_update_res_residual( self, update_list: list[paddle.Tensor], update_name: str = "g1" ) -> paddle.Tensor: - nitem = len(update_list) uu = update_list[0]🧰 Tools
🪛 Ruff (0.8.0)
290-290: Use
config.get("learning_rate_dict")instead ofconfig.get("learning_rate_dict", None)Replace
config.get("learning_rate_dict", None)withconfig.get("learning_rate_dict")(SIM910)
392-392: Use
config.get("learning_rate_dict")instead ofconfig.get("learning_rate_dict", None)Replace
config.get("learning_rate_dict", None)withconfig.get("learning_rate_dict")(SIM910)
493-493: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
636-636: Use a context manager for opening files
(SIM115)
646-646: Use a context manager for opening files
(SIM115)
683-686: Use ternary operator
pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lrinstead ofif-else-blockReplace
if-else-block withpref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr(SIM108)
739-739: Loop control variable
iinot used within loop bodyRename unused
iito_ii(B007)
863-863: Local variable
moduleis assigned to but never usedRemove assignment to unused variable
module(F841)
1048-1048: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
🪛 GitHub Check: CodeQL
[warning] 415-415: Unreachable code
This statement is unreachable.
[warning] 647-651: File is not always closed
File is opened but is not closed.
[warning] 657-657: File is not always closed
File is opened but is not closed.
[notice] 875-875: Unused local variable
Variable module is not used.deepmd/pd/model/descriptor/repformer_layer.py (3)
967-967: Remove unused local variableng2The variable
ng2is assigned but never used. Removing it will clean up the code.Apply this diff to remove the unused variable:
ng2 = g2.shape[-1]🧰 Tools
🪛 Ruff (0.8.0)
967-967: Local variable
ng2is assigned to but never usedRemove assignment to unused variable
ng2(F841)
🪛 GitHub Check: CodeQL
[notice] 967-967: Unused local variable
Variable ng2 is not used.
1157-1160: Simplify theif-elseblock using a ternary operatorYou can replace the
if-elseblock with a single line using a ternary operator to make the code more concise.Apply this diff to simplify the code:
if cal_gg1: gg1 = _make_nei_g1(g1_ext, nlist) else: gg1 = None + # Simplify the above using a ternary operator + gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None🧰 Tools
🪛 Ruff (0.8.0)
1157-1160: Use ternary operator
gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else Noneinstead ofif-else-blockReplace
if-else-block withgg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None(SIM108)
1278-1278: Remove unused local variablenitemThe variable
nitemis assigned but never used. Removing it will clean up the code.Apply this diff to remove the unused variable:
def list_update_res_residual( self, update_list: list[paddle.Tensor], update_name: str = "g1" ) -> paddle.Tensor: - nitem = len(update_list) uu = update_list[0]🧰 Tools
🪛 Ruff (0.8.0)
1278-1278: Local variable
nitemis assigned to but never usedRemove assignment to unused variable
nitem(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (82)
.github/workflows/test_cuda.yml(1 hunks).pre-commit-config.yaml(2 hunks)deepmd/__init__.py(1 hunks)deepmd/common.py(2 hunks)deepmd/dpmodel/__init__.py(1 hunks)deepmd/dpmodel/atomic_model/__init__.py(1 hunks)deepmd/dpmodel/atomic_model/dp_atomic_model.py(1 hunks)deepmd/dpmodel/atomic_model/linear_atomic_model.py(2 hunks)deepmd/dpmodel/atomic_model/make_base_atomic_model.py(1 hunks)deepmd/dpmodel/atomic_model/pairtab_atomic_model.py(2 hunks)deepmd/dpmodel/common.py(1 hunks)deepmd/dpmodel/descriptor/__init__.py(1 hunks)deepmd/dpmodel/fitting/__init__.py(1 hunks)deepmd/dpmodel/fitting/dipole_fitting.py(3 hunks)deepmd/dpmodel/fitting/dos_fitting.py(3 hunks)deepmd/dpmodel/fitting/ener_fitting.py(3 hunks)deepmd/dpmodel/fitting/general_fitting.py(9 hunks)deepmd/dpmodel/fitting/invar_fitting.py(3 hunks)deepmd/dpmodel/fitting/polarizability_fitting.py(4 hunks)deepmd/dpmodel/fitting/property_fitting.py(3 hunks)deepmd/dpmodel/model/__init__.py(1 hunks)deepmd/dpmodel/model/make_model.py(2 hunks)deepmd/dpmodel/utils/__init__.py(1 hunks)deepmd/entrypoints/test.py(4 hunks)deepmd/env.py(1 hunks)deepmd/infer/__init__.py(1 hunks)deepmd/infer/deep_eval.py(2 hunks)deepmd/infer/model_devi.py(1 hunks)deepmd/jax/descriptor/__init__.py(1 hunks)deepmd/jax/env.py(1 hunks)deepmd/jax/fitting/__init__.py(1 hunks)deepmd/jax/model/__init__.py(1 hunks)deepmd/pd/__init__.py(1 hunks)deepmd/pd/entrypoints/main.py(1 hunks)deepmd/pd/loss/ener.py(1 hunks)deepmd/pd/model/atomic_model/dp_atomic_model.py(1 hunks)deepmd/pd/model/descriptor/__init__.py(1 hunks)deepmd/pd/model/descriptor/dpa1.py(1 hunks)deepmd/pd/model/descriptor/dpa2.py(1 hunks)deepmd/pd/model/descriptor/env_mat.py(1 hunks)deepmd/pd/model/descriptor/repformer_layer.py(1 hunks)deepmd/pd/model/descriptor/repformers.py(1 hunks)deepmd/pd/model/descriptor/se_a.py(1 hunks)deepmd/pd/model/descriptor/se_atten.py(1 hunks)deepmd/pd/model/descriptor/se_t_tebd.py(1 hunks)deepmd/pd/model/model/__init__.py(1 hunks)deepmd/pd/model/model/ener_model.py(1 hunks)deepmd/pd/model/model/make_model.py(1 hunks)deepmd/pd/model/network/init.py(1 hunks)deepmd/pd/model/network/network.py(1 hunks)deepmd/pd/model/task/__init__.py(1 hunks)deepmd/pd/model/task/ener.py(1 hunks)deepmd/pd/model/task/fitting.py(1 hunks)deepmd/pd/model/task/invar_fitting.py(1 hunks)deepmd/pd/train/training.py(1 hunks)deepmd/pd/train/wrapper.py(1 hunks)deepmd/pd/utils/__init__.py(1 hunks)deepmd/pd/utils/dataloader.py(1 hunks)deepmd/pd/utils/decomp.py(1 hunks)deepmd/pd/utils/env.py(1 hunks)deepmd/pd/utils/exclude_mask.py(1 hunks)deepmd/pd/utils/nlist.py(1 hunks)deepmd/pd/utils/region.py(1 hunks)deepmd/pd/utils/utils.py(1 hunks)deepmd/pt/cxx_op.py(1 hunks)deepmd/pt/loss/__init__.py(1 hunks)deepmd/pt/model/atomic_model/__init__.py(1 hunks)deepmd/pt/model/atomic_model/dp_atomic_model.py(1 hunks)deepmd/pt/model/atomic_model/linear_atomic_model.py(2 hunks)deepmd/pt/model/atomic_model/pairtab_atomic_model.py(2 hunks)deepmd/pt/model/descriptor/__init__.py(1 hunks)deepmd/pt/model/descriptor/dpa2.py(1 hunks)deepmd/pt/model/model/__init__.py(3 hunks)deepmd/pt/model/model/make_model.py(1 hunks)deepmd/pt/model/network/init.py(1 hunks)deepmd/pt/model/task/__init__.py(1 hunks)deepmd/pt/model/task/dipole.py(4 hunks)deepmd/pt/model/task/dos.py(3 hunks)deepmd/pt/model/task/ener.py(3 hunks)deepmd/pt/model/task/fitting.py(11 hunks)deepmd/pt/model/task/invar_fitting.py(4 hunks)deepmd/pt/model/task/polarizability.py(5 hunks)
⛔ Files not processed due to max files limit (15)
- deepmd/pt/model/task/property.py
- deepmd/pt/train/training.py
- deepmd/pt/train/wrapper.py
- deepmd/pt/utils/init.py
- deepmd/pt/utils/dataloader.py
- deepmd/pt/utils/env.py
- deepmd/pt/utils/plugin.py
- deepmd/tf/init.py
- deepmd/tf/common.py
- deepmd/tf/descriptor/init.py
- deepmd/tf/entrypoints/init.py
- deepmd/tf/entrypoints/main.py
- deepmd/tf/env.py
- deepmd/tf/fit/init.py
- deepmd/tf/fit/dipole.py
✅ Files skipped from review due to trivial changes (2)
- deepmd/jax/fitting/init.py
- deepmd/jax/model/init.py
🚧 Files skipped from review as they are similar to previous changes (7)
- deepmd/pd/utils/init.py
- deepmd/dpmodel/model/make_model.py
- deepmd/pd/model/task/init.py
- .github/workflows/test_cuda.yml
- deepmd/pd/init.py
- deepmd/pd/model/descriptor/env_mat.py
- deepmd/pd/model/descriptor/init.py
🧰 Additional context used
🪛 Ruff (0.8.0)
deepmd/pd/model/descriptor/se_t_tebd.py
129-129: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
136-136: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
387-387: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
454-454: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
521-521: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
528-528: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
709-713: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
733-733: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
787-787: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/utils/region.py
27-27: Local variable e is assigned to but never used
Remove assignment to unused variable e
(F841)
deepmd/pd/utils/dataloader.py
54-54: Use capitalized environment variable FLAGS_CUDNN_DETERMINISTIC instead of FLAGS_cudnn_deterministic
Replace FLAGS_cudnn_deterministic with FLAGS_CUDNN_DETERMINISTIC
(SIM112)
86-86: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
285-290: Use a single if statement instead of nested if statements
(SIM102)
286-290: Use a single if statement instead of nested if statements
Combine if statements using and
(SIM102)
310-310: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/task/ener.py
40-40: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/entrypoints/main.py
118-118: Local variable f is assigned to but never used
Remove assignment to unused variable f
(F841)
500-503: Use ternary operator FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args instead of if-else-block
Replace if-else-block with FLAGS = parse_args(args=args) if not isinstance(args, argparse.Namespace) else args
(SIM108)
deepmd/pd/utils/exclude_mask.py
17-17: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
25-25: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
70-70: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
78-78: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/model/task/invar_fitting.py
93-93: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
104-104: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/train/training.py
290-290: Use config.get("learning_rate_dict") instead of config.get("learning_rate_dict", None)
Replace config.get("learning_rate_dict", None) with config.get("learning_rate_dict")
(SIM910)
392-392: Use config.get("learning_rate_dict") instead of config.get("learning_rate_dict", None)
Replace config.get("learning_rate_dict", None) with config.get("learning_rate_dict")
(SIM910)
493-493: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
636-636: Use a context manager for opening files
(SIM115)
646-646: Use a context manager for opening files
(SIM115)
683-686: Use ternary operator pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr instead of if-else-block
Replace if-else-block with pref_lr = _lr.start_lr if _step_id < self.warmup_steps else cur_lr
(SIM108)
739-739: Loop control variable ii not used within loop body
Rename unused ii to _ii
(B007)
863-863: Local variable module is assigned to but never used
Remove assignment to unused variable module
(F841)
1048-1048: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/task/fitting.py
132-132: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
143-143: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
251-251: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
328-328: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/descriptor/dpa2.py
90-90: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
499-499: Loop control variable ii not used within loop body
Rename unused ii to _ii
(B007)
684-684: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
deepmd/pd/model/descriptor/repformers.py
81-81: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
360-360: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
412-412: Local variable ng1 is assigned to but never used
Remove assignment to unused variable ng1
(F841)
435-435: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
548-552: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
deepmd/pd/model/descriptor/se_atten.py
61-61: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
78-78: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
375-379: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
399-399: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
477-477: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/model/descriptor/dpa1.py
218-218: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
230-230: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
534-534: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
633-633: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/model/descriptor/repformer_layer.py
967-967: Local variable ng2 is assigned to but never used
Remove assignment to unused variable ng2
(F841)
1157-1160: Use ternary operator gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None instead of if-else-block
Replace if-else-block with gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None
(SIM108)
1278-1278: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
deepmd/pd/model/descriptor/se_a.py
47-52: Use contextlib.suppress(ImportError) instead of try-except-pass
Replace with contextlib.suppress(ImportError)
(SIM105)
76-76: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
82-82: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
226-226: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
361-361: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
414-414: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
420-420: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
607-611: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
631-631: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
643-643: Loop control variable ll not used within loop body
Rename unused ll to _ll
(B007)
717-717: Loop control variable compress_data_ii not used within loop body
Rename unused compress_data_ii to _compress_data_ii
(B007)
717-717: Loop control variable compress_info_ii not used within loop body
Rename unused compress_info_ii to _compress_info_ii
(B007)
deepmd/pd/model/model/make_model.py
361-361: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
362-362: Use key not in dict instead of key not in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/model/network/init.py
56-56: Ambiguous variable name: l
(E741)
328-328: No explicit stacklevel keyword argument found
(B028)
378-378: No explicit stacklevel keyword argument found
(B028)
deepmd/pd/utils/nlist.py
386-386: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
🪛 GitHub Check: CodeQL
deepmd/pd/model/descriptor/se_t_tebd.py
[notice] 386-386: Unused local variable
Variable env_mat is not used.
[notice] 451-451: Unused local variable
Variable nall is not used.
[notice] 770-770: Unused local variable
Variable nall is not used.
deepmd/pd/utils/env.py
[notice] 72-72: Unused global variable
The global variable 'inter_nthreads' is not used.
[notice] 72-72: Unused global variable
The global variable 'intra_nthreads' is not used.
[notice] 75-76: Commented-out code
This comment appears to contain commented-out code.
deepmd/pd/utils/utils.py
[notice] 55-55: Statement has no effect
This statement has no effect.
deepmd/pd/train/training.py
[warning] 415-415: Unreachable code
This statement is unreachable.
[warning] 647-651: File is not always closed
File is opened but is not closed.
[warning] 657-657: File is not always closed
File is opened but is not closed.
[notice] 875-875: Unused local variable
Variable module is not used.
deepmd/pd/model/descriptor/dpa2.py
[warning] 663-663: Variable defined multiple times
This assignment to 'env_mat' is unnecessary as it is redefined before this value is used.
This assignment to 'env_mat' is unnecessary as it is redefined before this value is used.
[warning] 678-678: Variable defined multiple times
This assignment to 'env_mat' is unnecessary as it is redefined before this value is used.
[notice] 694-694: Unused local variable
Variable env_mat is not used.
deepmd/pd/model/descriptor/repformers.py
[notice] 416-416: Unused local variable
Variable ng1 is not used.
deepmd/pd/model/descriptor/se_atten.py
[notice] 431-431: Unused local variable
Variable nall is not used.
deepmd/pd/model/descriptor/dpa1.py
[notice] 529-529: Unused local variable
Variable env_mat is not used.
[notice] 600-600: Unused local variable
Variable nall is not used.
deepmd/pd/model/descriptor/repformer_layer.py
[notice] 967-967: Unused local variable
Variable ng2 is not used.
[notice] 1277-1277: Unused local variable
Variable nitem is not used.
deepmd/pd/model/descriptor/se_a.py
[notice] 50-50: Empty except
'except' clause does nothing but pass and there is no explanatory comment.
[notice] 323-323: Unused local variable
Variable env_mat is not used.
[warning] 655-655: Variable defined multiple times
This assignment to 'ti' is unnecessary as it is redefined before this value is used.
deepmd/pd/model/model/make_model.py
[notice] 323-328: Commented-out code
This comment appears to contain commented-out code.
deepmd/pd/utils/nlist.py
[notice] 395-395: Unused local variable
Variable nall is not used.
🔇 Additional comments (109)
deepmd/dpmodel/__init__.py (1)
27-28: LGTM! Improved organization of public interface
The reordering of items in __all__ follows a logical pattern, grouping related items together (constants, core classes, output definitions, operations, and utility functions) which improves code readability while maintaining the same functionality.
Also applies to: 30-31, 36-37
deepmd/pt/model/network/init.py (1)
133-134: LGTM! Improved type checking for negative_slope parameter.
The changes correctly handle the edge case where boolean values (which are instances of int in Python) could be incorrectly accepted as valid negative_slope values for leaky_relu activation. This fix ensures that only proper numeric types (int or float, but not bool) are accepted.
deepmd/pt/model/atomic_model/__init__.py (1)
50-50: LGTM: Addition of DPDipoleAtomicModel to all
The addition of DPDipoleAtomicModel to the public interface is consistent with its import statement above.
deepmd/pt/model/task/invar_fitting.py (2)
59-60: LGTM: Well-implemented case embedding support.
The addition of dim_case_embd parameter is well-documented and properly integrated:
- Clear docstring with type information
- Sensible default value for backward compatibility
- Correct propagation to parent class
Also applies to: 96-96, 119-119
Line range hint 183-184: LGTM: JIT compatibility fix.
The class-level type declaration for exclude_types is a good practice for ensuring JIT compatibility with PyTorch 2.0.0.
deepmd/jax/env.py (1)
21-21: LGTM! Verify usage of jax_export in codebase.
The addition of jax_export to the public API is appropriate as it provides access to JAX's function export capabilities.
Let's verify how jax_export is being used in the codebase:
✅ Verification successful
LGTM! The jax_export usage is verified and appropriate.
The codebase shows proper usage of jax_export in model serialization and deserialization:
- Used in
deepmd/jax/model/hlo.pyfor deserializing StableHLO models - Used in
deepmd/jax/utils/serialization.pyfor exporting functions with symbolic shapes
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check usage of jax_export to ensure it's being used appropriately
# and the public API exposure is necessary
# Search for imports of jax_export
rg "from deepmd.jax.env import .*jax_export.*" -A 2
# Search for direct usage of jax_export
rg "jax_export\." -A 2
Length of output: 1183
deepmd/infer/model_devi.py (2)
214-216: LGTM!
The f-string formatting for the additional columns is clean and consistent with modern Python practices.
Line range hint 1-24: Well-structured type hints and documentation
The code demonstrates excellent use of type hints, function overloads, and comprehensive documentation. The type hints will improve IDE support and help catch type-related issues early in development.
deepmd/dpmodel/fitting/ener_fitting.py (1)
35-35: LGTM: Clean implementation of case embedding dimension support
The new dim_case_embd parameter is properly integrated into both the class constructor and superclass initialization with appropriate typing and default value.
Also applies to: 59-59
deepmd/pt/cxx_op.py (1)
53-55: Verify consistent error handling across the codebase.
Let's ensure similar ABI compatibility checks are consistently implemented across other backend implementations.
✅ Verification successful
ABI compatibility checks are consistently implemented across PyTorch and TensorFlow backends
The verification shows that ABI compatibility checks are properly implemented in both backend implementations:
deepmd/pt/cxx_op.py: PyTorch backenddeepmd/tf/env.py: TensorFlow backend
Both implementations:
- Check CXX11_ABI_FLAG compatibility between the compiled package and runtime
- Raise similar, descriptive error messages when ABI mismatch is detected
- Include additional warning for RHEL6/7 when CXX11_ABI_FLAG=1
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for consistent ABI compatibility checks across backend implementations
# Search for CXX11_ABI_FLAG references
echo "Checking CXX11_ABI_FLAG handling across the codebase:"
rg -A 5 "CXX11_ABI_FLAG" --type py
# Search for similar compatibility error handling patterns
echo "Checking compatibility error handling patterns:"
rg -A 5 "compiled with.*but.*runtime" --type py
Length of output: 4028
.pre-commit-config.yaml (2)
63-63: LGTM: Minor version update for clang-format
The update from v19.1.3 to v19.1.4 is a patch release and should only contain bug fixes.
32-32: Verify compatibility with ruff v0.8.0
The update to ruff v0.8.0 might introduce breaking changes. Please ensure compatibility with the project's Python codebase.
Let's check the Python files that might be affected by ruff's rules:
✅ Verification successful
Ruff v0.8.0 update is safe to proceed
The update from v0.7.4 to v0.8.0 is safe because:
- The project has explicit ruff configuration in both pyproject.toml and .pre-commit-config.yaml with well-defined rules
- The configuration uses specific rule selections (E, W, F, D, UP, etc.) rather than blanket enforcement
- The project has proper exclusions for third-party code and specific file types
- Ruff is used with
--fixargument, which means it will automatically fix compatible issues
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check Python files that might be affected by ruff's rules
# Exclude files in source/3rdparty as per the config
# Find all Python files excluding source/3rdparty
fd -e py -e pyi -e ipynb --exclude 'source/3rdparty' . | wc -l
# Show sample of potentially affected files
fd -e py -e pyi -e ipynb --exclude 'source/3rdparty' . | head -n 5
Length of output: 250
Script:
#!/bin/bash
# Let's check the ruff configuration in the project to understand the rules being enforced
# Check for pyproject.toml which might contain ruff config
cat pyproject.toml 2>/dev/null || echo "No pyproject.toml found"
# Check for ruff.toml
cat ruff.toml 2>/dev/null || echo "No ruff.toml found"
# Check if there's any ruff config in setup.cfg
cat setup.cfg 2>/dev/null || echo "No setup.cfg found"
# Also check the pre-commit config for ruff settings
cat .pre-commit-config.yaml
Length of output: 19241
deepmd/dpmodel/model/__init__.py (1)
32-32: Verify the intentional exposure of DPModelCommon class
Adding DPModelCommon to __all__ makes it part of the public API. While this aligns with the DPA-2 support objectives, please ensure:
- This exposure is intentional and required for DPA-2 support
- The class is properly documented with its public interface
- The class maintains backward compatibility as it's now part of the public contract
Let's check the documentation and usage:
✅ Verification successful
Exposing DPModelCommon is appropriate and well-documented
The verification shows that DPModelCommon is:
- A well-established base class used across multiple model implementations (EnergyModel, PropertyModel, PolarModel, etc.)
- Already exposed in both PyTorch (
pt) and Paddle (pd) backend implementations - Properly documented with clear purpose: "A base class to implement common methods for all the Models"
- Used consistently in the codebase with proper inheritance patterns
The exposure aligns with the existing architecture where the class is already public in other backends, making it consistent across the codebase.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check if DPModelCommon is properly documented
rg -A 10 "class DPModelCommon"
# Check for existing usage patterns
rg "from deepmd\.dpmodel\.model import DPModelCommon"
Length of output: 2206
Script:
#!/bin/bash
# Let's check for any imports or usage of DPModelCommon in a broader context
rg "DPModelCommon" -B 2 -A 2
# Let's also check the docstring and class definition more thoroughly
ast-grep --pattern 'class DPModelCommon {
$$$
}'
Length of output: 15419
deepmd/env.py (1)
17-21: LGTM! Improved organization of exported symbols.
The reordering of __all__ entries improves readability by grouping related precision variables together. This change is purely cosmetic and doesn't affect functionality.
Let's verify if these exports are documented consistently:
✅ Verification successful
Based on the search results, I can generate the final response since I have enough information to verify the exports usage:
LGTM! The reordering of exports aligns with their usage patterns.
The reordering of __all__ entries is well-justified as these precision-related variables are heavily used throughout the codebase:
GLOBAL_NP_FLOAT_PRECISIONandGLOBAL_ENER_FLOAT_PRECISIONare used extensively in data processing, model training, and testing modulesglobal_float_precis used in configuration summaries- The precision variables are logically grouped together, improving code organization
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if the exports are documented in the module docstring
# or related documentation files.
# Search for documentation of these exports in Python files and docs
rg -t py -t rst "GLOBAL_ENER_FLOAT_PRECISION|GLOBAL_NP_FLOAT_PRECISION|global_float_prec|SHARED_LIB_MODULE" --glob '!env.py'
Length of output: 36320
deepmd/dpmodel/fitting/dipole_fitting.py (1)
134-134: LGTM!
The dim_case_embd parameter is correctly forwarded to the superclass constructor.
deepmd/infer/__init__.py (1)
Line range hint 19-32: Well-implemented factory function, verify usage patterns
The implementation looks good with proper documentation and type hints. Since this is a compatibility function, let's verify its usage across the codebase.
deepmd/entrypoints/test.py (5)
707-709: LGTM: String formatting update in test_dos function
The change to f-string formatting improves readability and follows modern Python best practices.
721-723: LGTM: String formatting update in test_dos function (atomic DOS)
The change to f-string formatting is consistent with the previous changes and maintains the same output format.
863-865: LGTM: String formatting update in test_property function
The change to f-string formatting follows the same pattern as in test_dos, maintaining consistency across the codebase.
877-879: LGTM: String formatting update in test_property function (atomic property)
The change to f-string formatting is consistent with all other changes in the file.
Line range hint 707-879: Verify output file format compatibility
While the string formatting changes are correct, we should verify that the output file format remains compatible with any downstream tools or scripts that might parse these files.
✅ Verification successful
Based on the documentation and code analysis, I can now provide a final response:
Output format is well-documented and used consistently
The output file format is explicitly documented in doc/model/train-fitting-dos.md and is used consistently in tests. The format consists of:
- For
*.dos.out.*: Matrix with shape(2, numb_dos)containing reference and predicted values - For
*.property.out.*: Similar two-column format with reference and predicted values - Files include a descriptive header:
{system} - {ii}: data_X pred_X
The changes in the code maintain this exact format, and the only test that reads these files (source/tests/pt/test_dp_test.py) expects this same structure. No downstream tools were found that would be impacted by these changes.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if there are any tools or scripts that parse the .dos.out or .property.out files
# Search for files that might process the test output files
rg -g '!*.py' -l "\.(?:dos|property)\.out\."
# Search for file processing patterns in Python files
rg -t py "open\(.*(?:dos|property)\.out\."
Length of output: 122
Script:
#!/bin/bash
# Let's check the content of the documentation to understand the expected format
cat doc/model/train-fitting-dos.md
# Let's also search for any code that reads these output files
rg -t py "read|load.*\.(?:dos|property)\.out"
# Let's check if there are any test files that might verify the output format
fd "test.*\.py$" | xargs rg "\.(?:dos|property)\.out"
Length of output: 23605
deepmd/pd/model/network/init.py (3)
1-7: LGTM: Proper licensing and attribution
The file correctly includes the LGPL-3.0-or-later license and properly acknowledges the PyTorch source code adaptation.
86-187: LGTM: Well-implemented utility functions
The utility functions are well-implemented with:
- Comprehensive gain calculations for different activation functions
- Proper error handling for invalid inputs
- Clear documentation with mathematical formulas
189-456: LGTM: Well-implemented initialization functions
The public initialization functions are well-implemented with:
- Comprehensive documentation including mathematical formulas
- Proper error handling
- Consistent API design across different initialization methods
- Clear examples in docstrings
🧰 Tools
🪛 Ruff (0.8.0)
328-328: No explicit stacklevel keyword argument found
(B028)
378-378: No explicit stacklevel keyword argument found
(B028)
deepmd/dpmodel/fitting/property_fitting.py (2)
103-103: LGTM: Parameter correctly passed to parent class.
The new dim_case_embd parameter is properly passed to the parent class constructor.
116-116: Verify version compatibility change impact.
The version compatibility check has been updated from 2.1 to 3.1, which suggests breaking changes in the data format. This change needs to be verified for:
- Backward compatibility with existing models
- Alignment with the DPA-2 support objectives mentioned in the PR
Let's check for other version-related changes:
✅ Verification successful
Let me gather more specific information about version compatibility and migration:
Let me try one more verification to check for version-related changes in Python files:
Based on the gathered information, I can now provide a final response:
Version compatibility change is aligned with DPA-2 support and maintains backward compatibility
The version change from 2.1 to 3.1 is part of a consistent update across the codebase:
- All DPA-2 related components use version 3.1 (deepmd/pd/model/descriptor/dpa2.py)
- The change is consistently applied across all fitting modules (ener, dipole, dos, property, etc.)
- The model compatibility documentation shows that DeePMD-kit maintains backward compatibility between versions
- The
dp convert-fromtool is available for converting older models to newer versions
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for other version compatibility checks to ensure consistency
rg -A 2 "check_version_compatibility.*3.*1"
# Look for migration guides or documentation updates
fd -e md -e rst | xargs rg -l "version|3.1|migration|upgrade"
Length of output: 5407
Script:
#!/bin/bash
# Check content of model compatibility documentation
cat doc/troubleshooting/model-compatability.md
# Check for any version-related changes in recent commits
git log --oneline -n 10 | grep -i "version\|3.1\|compatibility"
# Look for version constants or definitions
rg "VERSION|version" --type python -A 2
Length of output: 1093
Script:
#!/bin/bash
# Look for version constants or definitions in Python files
rg "VERSION|version" -t py -A 2
# Check for any recent changes in the version compatibility documentation
git log -p doc/troubleshooting/model-compatability.md | head -n 50
Length of output: 73661
deepmd/pd/utils/nlist.py (1)
1-521: Overall implementation looks good!
The neighbor list implementation is well-structured, properly documented, and handles edge cases appropriately. The suggested improvements are mainly for optimization and code cleanliness.
🧰 Tools
🪛 Ruff (0.8.0)
386-386: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
🪛 GitHub Check: CodeQL
[notice] 395-395: Unused local variable
Variable nall is not used.
deepmd/pt/model/task/dipole.py (3)
54-55: LGTM! Clean implementation of case embedding support.
The new parameter dim_case_embd is well-documented and properly integrated into both the constructor and parent class initialization.
Also applies to: 86-86, 109-109
135-135: Verify version compatibility impact.
The version compatibility check has been updated from version 2 to 3. This change might affect the deserialization of existing models.
Let's check for existing models that might be affected:
✅ Verification successful
Based on the search results, I can see that version 3 is consistently used across multiple components in both the PyTorch (pt) and PaddlePaddle (pd) implementations, including:
- deepmd/pt/model/task/dipole.py
- deepmd/pt/model/task/ener.py
- deepmd/pt/model/task/invar_fitting.py
- deepmd/pd/model/task/invar_fitting.py
- deepmd/pd/model/task/ener.py
The version 3 compatibility check appears to be a standardized update across the codebase rather than an isolated change. The version bump is consistently implemented with backward compatibility support (checking versions 1-3).
Version compatibility change is safe.
The version check allows models from versions 1-3, maintaining backward compatibility while standardizing on version 3 for new models. This is a safe change that won't break existing model deserialization.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Search for model files and version checks
# Look for model files and version references
rg -A 2 '"@version"'
Length of output: 79250
Line range hint 86-109: Ensure test coverage for case embedding functionality.
Since this adds support for case-specific embeddings, please ensure:
- Test cases cover both zero and non-zero
dim_case_embdscenarios - Integration tests verify compatibility with DPA-2 workflow
Let's check for existing tests:
✅ Verification successful
Let's examine the test content to verify the coverage of case embedding functionality.
Let's check the specific test content to understand the coverage better.
Based on the test files examined, I can now provide a final response:
Test coverage for case embedding functionality is present and comprehensive.
The codebase already includes thorough test coverage for dim_case_embd parameter:
- The universal test suite (
test_fitting.py) includes parameterized tests that verify both zero and non-zerodim_case_embdscenarios throughFittingParamDipoletest cases. - Integration tests in
test_dipole_fitting.pyverify:- Rotational invariance
- Permutation invariance
- Translational invariance
- Automatic differentiation
- Model inference
The test coverage is adequate as it verifies both the functionality and physical properties required for the dipole model with case embeddings.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Look for related test files
fd "test.*dipole"
Length of output: 467
Script:
#!/bin/bash
# Check test content for case embedding related tests
rg -l "dim_case_embd" source/tests/
rg "dim_case_embd.*=.*[1-9]" source/tests/ -A 3
Length of output: 192
Script:
#!/bin/bash
# Check content of relevant test files
cat source/tests/pt/model/water/multitask_sharefit.json
cat source/tests/universal/dpmodel/fitting/test_fitting.py
# Also check dipole-specific test files for any test cases
cat source/tests/pt/model/test_dipole_fitting.py
cat source/tests/consistent/model/test_dipole.py
Length of output: 27621
deepmd/pd/model/descriptor/repformers.py (1)
1-576: Overall well-structured implementation with minor improvements needed
The implementation is well-documented and follows a clear structure. However, there are several areas for improvement:
- Replace mutable default arguments
- Remove unused variables
- Implement missing communication dictionary handling
- Add parameter validation
- Simplify conditional logic
Once these improvements are addressed, the code will be more robust and maintainable.
🧰 Tools
🪛 Ruff (0.8.0)
81-81: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
360-360: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
412-412: Local variable ng1 is assigned to but never used
Remove assignment to unused variable ng1
(F841)
435-435: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
548-552: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
🪛 GitHub Check: CodeQL
[notice] 416-416: Unused local variable
Variable ng1 is not used.
deepmd/dpmodel/fitting/polarizability_fitting.py (2)
Line range hint 1-324: Implementation maintains mathematical correctness
The addition of case embedding support maintains backward compatibility and doesn't affect the core polarizability fitting functionality. The mathematical operations and rotational equivariance handling remain intact.
192-192: Verify version change documentation
The version bump from 3 to 4 is correctly implemented in both serialize and deserialize methods.
Let's check if this version change is documented:
Also applies to: 203-203
✅ Verification successful
Version change is properly documented
The version bump from 3 to 4 is related to the addition of dim_case_embd parameter, which is well documented in:
- Documentation at
doc/train/multi-task-training.mdexplaining its purpose for multi-task training - API documentation in
deepmd/utils/argcheck.pydescribing it as "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." - Example usage in
examples/water_multi_task/pytorch_example/input_torch_sharefit.json
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for version documentation in markdown files
rg -t md "version.*[3-4].*polar.*fitting"
# Search for other version constants or checks
rg "version.*[3-4].*polar.*fitting"
Length of output: 81
Script:
#!/bin/bash
# Let's try a broader search for version documentation
rg -t md "version.*[3-4]" --glob "doc/**"
rg -t md "version.*[3-4]" --glob "docs/**"
# Search for version-related changes in the polarizability fitting module
rg "version" deepmd/dpmodel/fitting/polarizability_fitting.py -A 2 -B 2
# Check for changelog or release notes
fd "CHANGELOG|RELEASE" -t f
Length of output: 4541
Script:
#!/bin/bash
# Let's check for recent changes or commits related to this version bump
git log -n 10 --oneline deepmd/dpmodel/fitting/polarizability_fitting.py
# Check for any documentation about dim_case_embd which triggered this version change
rg "dim_case_embd" -A 2 -B 2
# Look for any pull request description or commit message that might document this change
git log -n 5 --pretty=full deepmd/dpmodel/fitting/polarizability_fitting.py
Length of output: 41707
deepmd/dpmodel/fitting/invar_fitting.py (2)
159-159: LGTM: Case embedding integration
The dim_case_embd parameter is correctly passed to the superclass constructor, maintaining consistency with the parent class implementation.
188-188: Verify version compatibility change implications
The version compatibility check has been updated from 2 to 3, which indicates a breaking change. This change needs to be properly documented for users upgrading to the new version.
Let's check if this breaking change is documented:
deepmd/pt/model/descriptor/dpa2.py (2)
Line range hint 1-24: LGTM: Well-structured imports and licensing
The imports are well-organized into logical groups with proper type hints. The SPDX license header is correctly included.
Line range hint 302-350: Verify model compression safety
The enable_compression method performs critical checks but could miss some edge cases.
Let's verify the compression-related code:
✅ Verification successful
Based on the shell script results, I can now generate the final response:
Model compression is safely implemented
The codebase shows a consistent pattern for compression handling:
- All descriptors initialize
self.compress = Falsein their constructors - The
enable_compression()method is called only once during model compression - The compression flag is set to
Trueonly after all necessary compression setup is complete - The DPA2 descriptor properly sets up compression through its components
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check for potential compression-related issues in the codebase
# Look for places where compression might be enabled multiple times
rg -A 2 "enable_compression|compress\s*=\s*True"
# Check for potential race conditions in compression
ast-grep --pattern 'self.compress = $_'
Length of output: 21032
deepmd/pd/utils/env.py (3)
39-68: LGTM: Well-structured precision configuration
The precision configuration is comprehensive and includes proper validation checks. The mapping between numpy and paddle types is complete and well-organized.
153-169: LGTM: Complete and accurate exports
The __all__ list properly exports all necessary components of the module.
70-77:
Review threading configuration implementation
The current implementation has several issues:
- The variables
inter_nthreadsandintra_nthreadsare unused - There's commented out code that seems important for thread configuration
Let's verify the threading configuration in other files:
Consider either:
- Implementing the commented thread configuration
- Removing the unused thread configuration code
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 72-72: Unused global variable
The global variable 'inter_nthreads' is not used.
[notice] 72-72: Unused global variable
The global variable 'intra_nthreads' is not used.
[notice] 75-76: Commented-out code
This comment appears to contain commented-out code.
deepmd/dpmodel/common.py (1)
224-226: LGTM: Appropriate reordering of exports
The reordering of exports in __all__ improves readability while maintaining all necessary exports.
deepmd/pt/model/task/ener.py (2)
53-53: Verify integration of dim_case_embd parameter.
The addition of dim_case_embd parameter with a default value of 0 looks good. However, let's verify its integration with other components.
Also applies to: 71-71
✅ Verification successful
Let me gather more specific information about the integration by checking the actual implementation and usage patterns.
Let me check one more aspect to verify the inheritance chain and parameter passing:
Based on the results from our verification scripts, I can now provide the final response:
The dim_case_embd parameter integration is properly implemented
The parameter is consistently integrated across the codebase as evidenced by:
- Present in multiple fitting implementations (deepmd/pt/model/task/, deepmd/pd/model/task/, deepmd/tf/fit/*)
- Properly inherited through the fitting class hierarchy (EnergyFittingNet -> InvarFitting -> GeneralFitting)
- Default value of 0 maintains backward compatibility
- Implementation exists across all backends (PyTorch, TensorFlow, Paddle)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if dim_case_embd is consistently used across fitting classes
# and verify its integration with DPA-2 related code.
# Check other fitting classes for dim_case_embd parameter
echo "Checking fitting classes for dim_case_embd parameter..."
ast-grep --pattern 'class $_(InvarFitting) {
def __init__($$$, dim_case_embd: int = 0, $$$) {
$$$
}
}'
# Check for DPA-2 related usage
echo "Checking DPA-2 related usage..."
rg -l "dim_case_embd" --type py
Length of output: 1185
Script:
#!/bin/bash
# Check the implementation in base classes and usage patterns
echo "Checking InvarFitting base class implementation..."
rg -A 10 "class InvarFitting" deepmd/pt/model/task/invar_fitting.py
echo "\nChecking dim_case_embd usage in EnergyFittingNet..."
rg -A 10 "class EnergyFittingNet" deepmd/pt/model/task/ener.py
echo "\nChecking similar implementations in other fitting classes..."
rg -A 10 "def __init__.*dim_case_embd" deepmd/pt/model/task/
Length of output: 1440
Script:
#!/bin/bash
# Check the base class implementation and parameter handling
echo "Checking GeneralFitting class implementation..."
rg -A 15 "class GeneralFitting" deepmd/pt/model/task/fitting.py
echo "\nChecking dim_case_embd usage in InvarFitting..."
rg -A 15 "__init__.*dim_case_embd" deepmd/pt/model/task/invar_fitting.py
echo "\nChecking if dim_case_embd is used in forward methods..."
rg -B 2 -A 10 "def forward.*dim_case_embd" deepmd/pt/model/task/
Length of output: 1031
83-83: Verify version compatibility across the codebase.
The version check update from 2 to 3 looks good, maintaining backward compatibility to version 1.
✅ Verification successful
The search results show consistent version compatibility checks across the codebase, with all relevant files using the same version parameters (3 as max version, 1 as min version) in their check_version_compatibility calls. This includes various model tasks (energy, dipole, DOS), descriptors, and fitting modules across different frameworks (PT, TF, PD). Let me verify the implementation of the version check function itself to ensure it handles these versions correctly.
Let me try to locate the version compatibility check function using ripgrep instead.
Let me check the complete implementation of the version compatibility function to ensure it properly handles the version checks.
Let me get the complete function implementation including the actual version check logic.
Version compatibility checks are consistent and correctly implemented across the codebase
The verification shows that:
- All relevant files consistently use version 3 as maximum and version 1 as minimum in their compatibility checks
- The
check_version_compatibilityfunction indeepmd/utils/version.pyproperly validates that versions fall within the supported range - The implementation maintains backward compatibility by accepting versions from 1 to 3
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify version compatibility checks across fitting classes
# and check for any migration guides or documentation updates.
# Check version checks in other fitting classes
echo "Checking version compatibility in fitting classes..."
rg "check_version_compatibility.*3" --type py
# Look for version-related documentation or migration guides
echo "Checking for version documentation..."
fd -e md -e rst | xargs rg -l "version.*3.*compatibility|migration"
Length of output: 2101
Script:
#!/bin/bash
# Check the implementation of check_version_compatibility function
ast-grep --pattern 'def check_version_compatibility'
Length of output: 54
Script:
#!/bin/bash
# Find the implementation of check_version_compatibility function
rg "def check_version_compatibility" -A 5
Length of output: 412
Script:
#!/bin/bash
# Get the complete implementation of check_version_compatibility function
rg "def check_version_compatibility" -A 15 deepmd/utils/version.py
Length of output: 554
Script:
#!/bin/bash
# Get the complete implementation including the version check logic
cat deepmd/utils/version.py
Length of output: 972
deepmd/pt/model/task/dos.py (3)
Line range hint 1-124: Implementation looks solid
The DOSFittingNet implementation is well-structured with proper type hints, clear serialization/deserialization logic, and good integration of the new features. The changes align well with the PR's objective of supporting DPA-2.
104-104: Verify version compatibility impact
The version compatibility check has been updated from 2 to 3. This change requires careful consideration of backward compatibility.
✅ Verification successful
Based on the search results, I can see that version compatibility checks are used consistently across the codebase, with most components using version 3 as the maximum supported version for their serialization/deserialization. Looking at the specific file in question (deepmd/pt/model/task/dos.py), the version check is:
check_version_compatibility(data.pop("@version", 1), 3, 1)This is consistent with other task modules in the codebase that also use version 3 as their maximum supported version, such as:
- deepmd/pt/model/task/ener.py
- deepmd/pt/model/task/property.py
- deepmd/pt/model/task/invar_fitting.py
- deepmd/pt/model/task/dipole.py
The version compatibility check ensures backward compatibility by accepting versions between 1 (minimum) and 3 (maximum), which is the standard pattern used across the codebase.
Version compatibility is correctly implemented
The version compatibility check in the DOS task module follows the established pattern used across other task modules in the codebase, maintaining consistency in version handling.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check version handling across the codebase
# Find all version compatibility checks
rg "check_version_compatibility" --type py -A 2
# Look for version-related migration guides or documentation
fd -e md -e rst | xargs rg -l "version|compatibility|migration"
Length of output: 32192
50-50: Verify the integration of dim_case_embd parameter
The addition of dim_case_embd parameter appears to be part of a broader enhancement for case-specific embeddings across fitting models. While the implementation looks correct, we should verify its integration with the DPA-2 support.
Also applies to: 77-77
✅ Verification successful
The dim_case_embd parameter is correctly integrated
The implementation of dim_case_embd in the DOS fitting model is consistent with other fitting models in the codebase. Key findings:
- The parameter is properly integrated into the base
Fittingclass and consistently passed through all fitting models including DOS, Energy, Dipole, and Polarizability - It's used for case-specific embeddings in multitask models with proper validation in training code
- The parameter is supported in PyTorch and Paddle backends (with explicit checks in TensorFlow backend to raise errors)
- Documentation indicates it's used for the number of model branches in multitask scenarios
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify consistent implementation of dim_case_embd across fitting models
# and its usage in DPA-2 related code
# Check if dim_case_embd is consistently implemented across fitting models
ast-grep --pattern 'class $_(InvarFitting) {
$$$
def __init__($$$, dim_case_embd: int = 0, $$$) {
$$$
}
$$$
}'
# Look for DPA-2 related usage
rg -l "dpa2|DPA2" --type py
Length of output: 2731
Script:
#!/bin/bash
# Let's check the implementation of dim_case_embd in fitting models and its usage in DPA2
# Check fitting models that inherit from InvarFitting
ast-grep --pattern 'class $name(InvarFitting)'
# Check how dim_case_embd is used in DPA2 implementation
rg -A 5 "dim_case_embd" --type py
# Check the base class InvarFitting to understand the parameter usage
ast-grep --pattern 'class InvarFitting'
Length of output: 47554
deepmd/pt/model/task/polarizability.py (4)
56-57: LGTM: Documentation is clear and consistent
The documentation for the new parameter dim_case_embd follows the existing style and provides a clear description.
134-134: LGTM: Parameter correctly forwarded to parent class
The dim_case_embd parameter is properly passed to the parent class constructor.
90-90: Verify the default value's impact
The addition of dim_case_embd with a default value of 0 maintains backward compatibility. However, let's verify its usage across the codebase.
✅ Verification successful
Default value of 0 for dim_case_embd is safe and consistent
The parameter dim_case_embd is consistently defined across the codebase with a default value of 0. The documentation indicates it's used for multitask models with case embedding, where it should match the number of model branches. The default value of 0 effectively disables this feature, maintaining backward compatibility.
Key findings:
- In TensorFlow implementation, non-zero values are explicitly rejected with error messages
- Documentation consistently describes its purpose for multitask models
- All fitting implementations follow the same pattern of defaulting to 0
- No instances found of non-zero usage in instantiations
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check usage patterns of dim_case_embd across the codebase
# Look for other fitting networks using this parameter
ast-grep --pattern 'class $_(GeneralFitting):
$$$
def __init__($$$, dim_case_embd: int = 0, $$$):
$$$'
# Look for actual instantiations with non-zero dim_case_embd
rg -A 2 'dim_case_embd\s*='
Length of output: 20428
198-198: Verify version change coordination
The version bump from 3 to 4 needs to be coordinated across related files to maintain consistency in serialization/deserialization.
Also applies to: 209-209
✅ Verification successful
Version change is properly coordinated
The version bump from 3 to 4 in the polarizability fitting module is consistent with the codebase. The search results show that:
- The polarizability fitting module (
deepmd/tf/fit/polar.pyanddeepmd/pt/model/task/polarizability.py) consistently uses version 4 and checks compatibility with versions 1 through 4 - Other modules maintain their own versioning based on their specific needs
- The version compatibility checks are properly implemented across the codebase using the
check_version_compatibilityutility
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check version numbers in related fitting network files
# Look for version numbers in serialize methods
rg -B 2 '@version"\s*=\s*[0-9]+' --type py
# Look for version compatibility checks
rg -B 2 'check_version_compatibility.*[0-9]+\s*,\s*[0-9]+\s*,\s*[0-9]+' --type py
Length of output: 128
Script:
#!/bin/bash
# Search for version-related patterns in Python files more broadly
rg -B 2 -A 2 '@version' --type py
# Also check for version compatibility checks with a broader pattern
rg -B 2 -A 2 'version.*compatibility' --type py --ignore-case
# Look for serialize/deserialize methods that might contain version info
rg -B 2 -A 2 'def (de)?serialize' --type py
Length of output: 113982
deepmd/common.py (3)
38-39: LGTM! Export list updated correctly.
The additions to __all__ properly expose the new type-safe constants and maintain alphabetical ordering.
Also applies to: 42-44
66-68: LGTM! Type hints properly guarded.
The type-related exports are correctly guarded by TYPE_CHECKING, ensuring type hints are available for static type checkers without runtime overhead.
Line range hint 47-57: Type literals and sets are well-defined.
The implementation provides proper type safety using Literal types for both precision and activation functions. The sets are correctly initialized using get_args().
Let's verify that all activation functions are properly implemented:
Also applies to: 59-60
✅ Verification successful
All activation functions are properly implemented and supported
The verification shows that all activation functions listed in _ACTIVATION are properly implemented and supported across the codebase:
-
TensorFlow implementations are provided through
tf.nnfunctions indeepmd/tf/common.py:- relu:
tf.nn.relu - relu6:
tf.nn.relu6 - softplus:
tf.nn.softplus - sigmoid:
tf.sigmoid - tanh: Custom implementation in
deepmd/tf/nvnmd/utils/network.py - gelu/gelu_tf: Custom implementations in
deepmd/tf/common.py
- relu:
-
PyTorch/Paddle implementations are provided in
deepmd/pt/utils/utils.pyanddeepmd/pd/utils/utils.py -
A unified activation function factory is implemented via
get_activation_func/get_activation_fnacross frameworks -
All activation functions are properly mapped in the tabulation utilities
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify implementation of all activation functions
# Search for implementation of each activation function listed in _ACTIVATION
echo "Searching for activation function implementations..."
for activation in relu relu6 softplus sigmoid tanh gelu gelu_tf; do
echo "Checking $activation implementation:"
rg -l "def.*${activation}.*\(" --type python
done
Length of output: 1575
Script:
#!/bin/bash
# Let's try a different approach without relying on file type detection
# Search for activation function implementations using plain text search
echo "Searching for activation function implementations..."
for activation in relu relu6 softplus sigmoid tanh gelu gelu_tf; do
echo "Checking $activation implementation:"
rg "def.*${activation}" --type-add 'py:*.py' --type py
done
# Also search for potential imports of these functions
echo -e "\nChecking imports of activation functions..."
rg "import.*activation" --type-add 'py:*.py' --type py
Length of output: 3598
Script:
#!/bin/bash
# Let's check for the remaining activation functions in imports and function calls
# since we didn't find direct implementations of relu, relu6, softplus, and sigmoid
echo "Checking for activation function imports and usage..."
rg "(relu|relu6|softplus|sigmoid)" --type-add 'py:*.py' --type py
# Also check for activation function getter/factory
echo -e "\nChecking activation function factory..."
rg "get_activation" --type-add 'py:*.py' --type py -A 5
Length of output: 19785
deepmd/dpmodel/utils/__init__.py (1)
39-61: LGTM! Well-organized exports list
The __all__ list is properly updated to include the new mask classes and maintains alphabetical ordering. All exported symbols correspond to their imports.
deepmd/pd/utils/exclude_mask.py (2)
41-61: LGTM! Well-implemented forward method
The forward method is well-documented and efficiently implemented, with proper shape handling and device placement.
104-156: Verify tensor operations in forward method
The implementation looks correct but let's verify the tensor operations:
The forward method implementation is robust with:
- Proper handling of virtual atoms
- Efficient tensor operations
- Device placement consistency
- Type conversion for the final mask
✅ Verification successful
Based on the search results, I can generate the final response as there is sufficient evidence to verify the tensor operations in the forward method:
Tensor operations in forward method are correctly implemented
The implementation follows established patterns found across the codebase:
paddle.take_along_axisusage matches similar implementations indeepmd/pd/utils/nlist.pyand other descriptor modules- Tensor reshaping operations align with common patterns in the codebase
- Device placement and type conversion are handled consistently with other modules
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Search for similar tensor operations in the codebase for validation
# Look for similar patterns of paddle.take_along_axis and tensor reshaping
# Search for similar tensor operations
rg -A 3 "take_along_axis|paddle\.where.*reshape|paddle\.concat.*axis"
# Search for similar mask computations
ast-grep --pattern 'mask = $_.reshape($$$).to($_.place)'
Length of output: 25814
deepmd/pd/utils/utils.py (1)
1-198: Verify integration with DPA-2 implementation
The utility functions look well-implemented for supporting DPA-2 in the Paddle backend. Let's verify the usage of these utilities in the DPA-2 implementation.
✅ Verification successful
Utilities are properly integrated with DPA-2 implementation
The verification shows that the utility functions are correctly integrated and used in the codebase:
to_numpy_arrayis used in DPA-2 implementation for handling model variables (davg, dstd)ActivationFnis properly integrated in the MLP network components that can be used by DPA-2- The tensor conversion utilities provide the necessary infrastructure for data type handling
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for usage of ActivationFn and tensor conversion utilities in DPA-2 related files
rg -l "ActivationFn|to_paddle_tensor|to_numpy_array" --type python
Length of output: 99
Script:
#!/bin/bash
# Search for usage of ActivationFn and tensor conversion utilities without file type restriction
rg -l "ActivationFn|to_paddle_tensor|to_numpy_array"
Length of output: 3655
Script:
#!/bin/bash
# Check specific usage in DPA-2 implementation
rg -C 2 "ActivationFn|to_paddle_tensor|to_numpy_array" deepmd/pd/model/descriptor/dpa2.py
Length of output: 941
Script:
#!/bin/bash
# Check usage of ActivationFn in network components that might be used by DPA-2
rg -C 3 "ActivationFn" deepmd/pd/model/network/
Length of output: 918
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 55-55: Statement has no effect
This statement has no effect.
deepmd/pd/model/model/make_model.py (3)
42-82: LGTM! Well-structured factory pattern implementation.
The factory function and class structure demonstrate good OOP principles with proper handling of dependencies and initialization.
479-595: LGTM! Well-documented and properly implemented helper methods.
The helper methods demonstrate:
- Clear separation of concerns
- Consistent error handling
- Comprehensive documentation with type hints
322-328: Verify JIT compatibility impact.
The commented type checking code was removed for JIT compatibility. Consider documenting this limitation in the code comments for future maintenance.
✅ Verification successful
Let me gather more context about the JIT compatibility and type checking patterns.
Let me check how the type conversion is actually handled after removing the type checking code.
JIT compatibility change is properly handled
The type checking code was correctly commented out for JIT compatibility, and the implementation properly handles type conversion by directly casting input tensors to match the coordinate dtype. The code now:
- Uses a list comprehension to convert all inputs to the same dtype as coordinates
- Maintains consistency across different model implementations (pd, pt, dpmodel)
- Avoids JIT-incompatible type checking while still ensuring type consistency
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check for similar JIT-related type checking patterns in the codebase
rg -A 2 "type checking.*would not pass jit"
Length of output: 850
Script:
#!/bin/bash
# Check for any other JIT-related comments or patterns
rg -B 2 -A 2 "jit" --type python
# Look for any documentation about JIT limitations
fd README.md | xargs rg -i "jit"
# Check the full context of the make_model files
fd make_model.py | xargs rg -B 5 -A 5 "type checking.*would not pass jit"
Length of output: 2929
Script:
#!/bin/bash
# Check the implementation that replaced the type checking
ast-grep --pattern 'def _convert_input_dtype($$$)'
# Look for any related type conversion methods
rg -B 3 -A 3 "astype\(coord\.dtype\)"
Length of output: 1212
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 323-328: Commented-out code
This comment appears to contain commented-out code.
deepmd/dpmodel/atomic_model/make_base_atomic_model.py (1)
71-77: LGTM!
The abstract method is well-defined with clear documentation and follows the established pattern of the class.
deepmd/pt/model/atomic_model/dp_atomic_model.py (1)
96-101: Add error handling and verify framework consistency.
- Consider adding error handling for invalid case indices.
- The implementation uses
fitting_netwhile the other framework usesfitting. Although this is correct due to framework differences, we should verify the behavior is consistent.
def set_case_embd(self, case_idx: int):
"""
Set the case embedding of this atomic model by the given case_idx,
typically concatenated with the output of the descriptor and fed into the fitting net.
"""
+ if not isinstance(case_idx, int) or case_idx < 0:
+ raise ValueError(f"case_idx must be a non-negative integer, got {case_idx}")
self.fitting_net.set_case_embd(case_idx)deepmd/dpmodel/atomic_model/pairtab_atomic_model.py (2)
123-130: LGTM: Clear implementation of unsupported feature
The implementation correctly raises NotImplementedError with a clear message indicating that case identification is not supported for this model type.
438-446: LGTM: Clear implementation of unsupported feature
The implementation correctly indicates that compression is not supported for this model type through a no-op implementation.
deepmd/dpmodel/atomic_model/linear_atomic_model.py (2)
137-143: LGTM: Proper propagation of case embedding
The implementation correctly propagates the case embedding to all sub-models, which is consistent with the linear model's role of combining multiple models.
439-445: LGTM: Appropriate handling of case embedding for ZBL model
The implementation correctly sets the case embedding only for the first model (dpmodel), which is appropriate as the ZBL component doesn't support case embedding.
deepmd/pt/model/atomic_model/pairtab_atomic_model.py (2)
144-151: LGTM: Consistent implementation with dpmodel version
The implementation correctly raises NotImplementedError and matches the behavior in the dpmodel version.
497-506: LGTM: Consistent implementation with dpmodel version
The implementation correctly indicates that compression is not supported and matches the behavior in the dpmodel version.
deepmd/pt/model/atomic_model/linear_atomic_model.py (2)
161-168: LGTM: Case embedding propagation is correctly implemented
The implementation properly propagates the case embedding index to all sub-models in a clean and straightforward manner.
572-579: LGTM: Case embedding correctly limited to DP model
The implementation correctly sets the case embedding only for the first model (dp_model), which is the expected behavior since ZBL model doesn't need case embedding.
Let's verify that the first model is indeed always the DP model:
✅ Verification successful
LGTM: Case embedding correctly limited to DP model
The implementation is verified to be correct. The __init__ method of DPZBLLinearEnergyAtomicModel explicitly takes dp_model as its first parameter and initializes the models list as [dp_model, zbl_model]. This confirms that self.models[0] in set_case_embd always refers to the DP model, making it correct to only set the case embedding for this model.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify that the first model in DPZBLLinearEnergyAtomicModel is always the DP model
ast-grep --pattern 'class DPZBLLinearEnergyAtomicModel {
$$$
def __init__($_, dp_model: DPAtomicModel, zbl_model: PairTabAtomicModel, $$$) {
$$$
models = [dp_model, zbl_model]
$$$
}
$$$
}'
Length of output: 210
Script:
#!/bin/bash
# Let's try a simpler pattern to find the class initialization
ast-grep --pattern 'class DPZBLLinearEnergyAtomicModel'
# Also search for the models initialization
rg "models = \[" deepmd/pt/model/atomic_model/linear_atomic_model.py -A 2 -B 2
Length of output: 724
Script:
#!/bin/bash
# Let's look at the full class definition to understand the model initialization
rg "class DPZBLLinearEnergyAtomicModel" deepmd/pt/model/atomic_model/linear_atomic_model.py -A 20
Length of output: 898
Script:
#!/bin/bash
# Let's look at the __init__ method implementation
rg "def __init__" deepmd/pt/model/atomic_model/linear_atomic_model.py -A 10 -B 2 --context-separator="---"
Length of output: 746
deepmd/dpmodel/fitting/dos_fitting.py (3)
64-64: LGTM: Parameter correctly forwarded to superclass.
The dim_case_embd parameter is properly passed to the parent class constructor.
Line range hint 94-95: Verify deserialization handling for bias_atom_e.
While the serialization of bias_atom_e has been added, please ensure that:
- The deserialization logic properly handles this new field
- Backward compatibility is maintained for older serialized models
✅ Verification successful
Deserialization of bias_atom_e is properly handled
The verification shows that bias_atom_e is properly handled in both serialization and deserialization:
- The base class
GeneralFittinghandles deserialization through itsdeserializemethod which properly restores variables from the@variablesdictionary wherebias_atom_eis stored - The class has proper initialization support for
bias_atom_ein its__init__method - Backward compatibility is maintained since
bias_atom_eis an optional parameter with a default value ofNone
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for deserialization handling of bias_atom_e
ast-grep --pattern 'def deserialize($$$):
$$$
bias_atom_e = $$$
$$$'
Length of output: 76
Script:
#!/bin/bash
# Let's first check for any deserialization method in the dos_fitting.py file
rg -A 10 "def deserialize" deepmd/dpmodel/fitting/dos_fitting.py
# Also search for any load or from_dict methods that might handle deserialization
rg -A 10 "def (load|from_dict)" deepmd/dpmodel/fitting/dos_fitting.py
# Let's also check if there are any other files that might handle deserialization for this class
rg "class DOSFitting"
Length of output: 965
Script:
#!/bin/bash
# Let's check the parent class InvarFitting for deserialization handling
rg -A 15 "def deserialize" deepmd/dpmodel/fitting/invar_fitting.py
# Also check if bias_atom_e is handled in the parent class initialization
rg -B 5 -A 10 "def __init__" deepmd/dpmodel/fitting/invar_fitting.py
# Let's also check if there are any references to bias_atom_e in the parent class
rg "bias_atom_e" deepmd/dpmodel/fitting/invar_fitting.py
Length of output: 1150
Script:
#!/bin/bash
# Let's check the base class GeneralFitting for deserialization handling
rg -A 15 "def deserialize" deepmd/dpmodel/fitting/general_fitting.py
# Also check if bias_atom_e is handled in the base class initialization
rg -B 5 -A 10 "def __init__" deepmd/dpmodel/fitting/general_fitting.py
# Let's also look for any bias_atom_e references in the base class
rg "bias_atom_e" deepmd/dpmodel/fitting/general_fitting.py
Length of output: 2094
39-39: Verify the usage of dim_case_embd parameter.
The addition of dim_case_embd parameter looks good. Let's verify its consistent usage across other fitting networks.
✅ Verification successful
The dim_case_embd parameter is consistently implemented across the codebase
The parameter is properly integrated across the codebase with consistent implementation:
- It's used in all fitting networks (DOS, Energy, Dipole, Polarizability, etc.) with the same default value of 0
- Documentation clearly states its purpose: "The dimension of case embedding. When training or fine-tuning a multitask model with case embeddings, this number should be set to the number of model branches."
- Implementation details show:
- Proper initialization in base classes (general_fitting.py)
- Consistent handling in both PyTorch and Paddle backends
- TensorFlow backend explicitly raises ValueError when dim_case_embd > 0
- Test cases and examples (water_multi_task) demonstrate its usage in multitask scenarios
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for other fitting networks to ensure consistent parameter usage
ast-grep --pattern 'class $_FittingNet(InvarFitting):
def __init__(self, $$$, dim_case_embd: int = 0, $$$):
$$$'
Length of output: 120
Script:
#!/bin/bash
# Let's try a broader search for fitting networks and their initialization
ast-grep --pattern 'class $_FittingNet($_):
def __init__($$$):
$$$'
# Also search for direct references to dim_case_embd
rg "dim_case_embd" -A 2
# Check the parent class InvarFitting
ast-grep --pattern 'class InvarFitting($_):
def __init__($$$):
$$$'
Length of output: 26765
deepmd/infer/deep_eval.py (1)
Line range hint 505-519: Verify error handling with unit tests
The error handling in _standard_input is critical for the DPA-2 support. Consider adding unit tests to verify these error cases:
- Wrong fparam dimensions
- Wrong aparam dimensions
- Mixed combinations of dimensions
Would you like me to help generate comprehensive unit tests for these error cases?
deepmd/pd/model/descriptor/dpa2.py (1)
610-614: Ensure compatibility checks are clear and updated
The deserialize method performs a version compatibility check. Confirm that version numbers and error messages are informative and up-to-date.
Please verify that the check_version_compatibility function correctly handles version differences and provides meaningful messages to the user.
deepmd/pt/model/task/__init__.py (1)
32-32: LGTM: Appropriate addition to module exports
The addition of BaseFitting to __all__ is consistent with the module's purpose and the PR's objective of enhancing the fitting framework.
deepmd/pt/model/descriptor/__init__.py (2)
53-58: LGTM: Comprehensive descriptor API expansion
The addition of these descriptor classes (DescrptSeA, DescrptSeAttenV2, DescrptSeR, DescrptSeT, DescrptSeTTebd) enriches the module's public API and aligns with the PR's objective of supporting DPA-2.
46-46: Verify removal of DescrptBlockRepformers from exports
The removal of this class from exports might impact existing code that depends on it.
deepmd/pt/model/model/__init__.py (2)
Line range hint 199-209: LGTM: Improved model instantiation and configuration tracking
The changes enhance the get_zbl_model function by:
- Using proper variable assignment for better debugging
- Storing the model definition script, which aids in model reproduction and debugging
300-312: Verify impact of model class removals from exports
Several model classes have been removed from exports (DOSModel, DPZBLModel, EnergyModel, SpinModel) while others have been added (LinearEnergyModel, PolarModel). This significant API change requires verification.
deepmd/pd/train/wrapper.py (2)
1-20: LGTM! Well-structured imports and type definitions.
The file header is properly licensed, imports are well-organized, and type hints are appropriately used.
67-73: Verify implementation timeline for parameter sharing.
The share_params method is critical for DPA-2 support but is currently not implemented. This might block the functionality described in the PR objectives.
Would you like help implementing the parameter sharing functionality?
deepmd/pd/utils/decomp.py (2)
1-21: Well-documented file purpose and clear API exports.
The file header clearly explains its purpose and temporary nature. The imports are minimal and the __all__ list properly defines the public API.
137-139: LGTM! Clear and well-documented aliases.
The aliases provide convenient access to the decomposed functions while maintaining clear documentation of their purpose.
deepmd/pd/model/task/invar_fitting.py (3)
1-33: LGTM: Well-organized imports and setup
The imports are properly organized, typed, and follow best practices. The global configuration and logging setup are appropriate.
35-85: Consider removing or implementing the unsupported parameter
The dim_case_embd parameter is documented as "(Not supported yet)". Consider either:
- Implementing the functionality
- Removing the parameter if it's not planned for immediate implementation
- Adding a TODO comment with more details about the implementation plan
This will help prevent confusion for users of the API.
Let's check if this parameter is used in other parts of the codebase:
150-150: Verify version compatibility requirements
The version compatibility check requires version 3 or lower. Please ensure this aligns with the project's versioning strategy and document any breaking changes that necessitated this version requirement.
Let's check the version compatibility usage across the codebase:
✅ Verification successful
Version compatibility check is consistent with project versioning
Based on the evidence gathered:
- The version check requiring version 3 or lower is consistent with DeePMD-kit's versioning strategy, as seen in the project's version history in README.md where v3 is the latest major version.
- The
check_version_compatibilityfunction is consistently used across multiple components with version 3 as the maximum supported version, as shown by the widespread usage in the codebase. - The version check is properly documented and follows a standard pattern of ensuring backward compatibility while preventing usage of unsupported future versions.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for version compatibility checks
rg "check_version_compatibility.*3" --type python
Length of output: 82
Script:
#!/bin/bash
# Search for version compatibility checks without file type restriction
rg "check_version_compatibility.*3"
# Also search for the function definition to understand its implementation
ast-grep --pattern 'def check_version_compatibility($$$)'
# Look for other version-related constants or patterns
rg "@version"
Length of output: 39892
Script:
#!/bin/bash
# Let's look for the implementation of check_version_compatibility function
ast-grep --pattern 'def check_version_compatibility($$$)'
# Also search for any version-related documentation
rg -l "version" README.md docs/
Length of output: 147
Script:
#!/bin/bash
# Search for check_version_compatibility function with ripgrep
rg "def check_version_compatibility" -A 5
# Look at the README content for version information
cat README.md
Length of output: 9072
Script:
#!/bin/bash
# Get the full implementation of check_version_compatibility
rg "def check_version_compatibility" -A 15 deepmd/utils/version.py
Length of output: 554
deepmd/pd/model/model/__init__.py (3)
1-41: LGTM! Well-structured module setup with clear documentation.
The module documentation and import organization provide a clear understanding of the module's purpose and dependencies.
137-144: LGTM! Well-defined module exports.
The __all__ list clearly defines the public API of the module, including all necessary models and functions.
1-144: Verify DPA-2 compatibility across the codebase.
As this implementation is part of DPA-2 support in the Paddle backend, we should verify its integration with other components.
Let's check for DPA-2 related implementations and dependencies:
✅ Verification successful
Let me analyze the DPA-2 related files and implementations to verify compatibility. I'll generate another script to check specific DPA-2 implementations:
Based on the extensive search results, I can now provide a definitive response about the DPA-2 compatibility:
DPA-2 implementation is properly integrated across the codebase
The model implementation in deepmd/pd/model/model/__init__.py is compatible with DPA-2 because:
- The descriptor system is properly integrated through the
BaseDescriptorabstraction - DPA-2 implementations exist across all backends:
- Paddle backend:
deepmd/pd/model/descriptor/dpa2.py - PyTorch backend:
deepmd/pt/model/descriptor/dpa2.py - JAX backend:
deepmd/jax/descriptor/dpa2.py
- Paddle backend:
- Comprehensive test coverage exists for DPA-2:
- Energy model tests
- Force model tests
- Training tests
- Multi-task tests
- Consistency tests across backends
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check for DPA-2 related implementations and potential integration points
# Search for DPA-2 related code
echo "Searching for DPA-2 related implementations..."
rg -l "dpa2|DPA-2|DPA2" --type py
# Check for model type definitions that might need DPA-2 support
echo "Checking model type definitions..."
rg -A 5 "class.*Model\(" --type py
# Look for potential integration points in descriptor implementations
echo "Checking descriptor implementations..."
rg -A 5 "class.*Descriptor\(" --type py
Length of output: 52391
Script:
#!/bin/bash
# Check DPA-2 descriptor implementations and their integration points
# Look for DPA-2 descriptor class definitions
echo "Checking DPA-2 descriptor implementations..."
rg -A 10 "class.*DPA2" --type py
# Check for DPA-2 imports and usage
echo "Checking DPA-2 imports and usage..."
rg "from.*dpa2.*import|import.*dpa2" --type py
# Look for DPA-2 specific configurations
echo "Looking for DPA-2 configurations..."
rg "dpa2.*config|DPA2.*config" --type py
Length of output: 56382
deepmd/pd/model/task/ener.py (2)
34-35: LGTM: Class registration and inheritance are well-defined.
The class is properly registered as an energy fitting network and inherits from the appropriate base class.
72-86: Verify version compatibility changes and migration path.
The version compatibility check has been updated from 2 to 3, which suggests breaking changes. Please ensure:
- These changes are documented
- A migration path exists for users
- The changes are mentioned in the changelog
✅ Verification successful
Let me check if there are any version-related changes in the codebase or any documentation in other locations.
Let me check the version compatibility implementation and documentation.
Let me check for any recent changes in version compatibility and related documentation.
Let me check the version compatibility tests and any recent changes to the version handling.
Based on the gathered information, I can now provide a final response:
Version compatibility changes are properly documented and supported
The version bump from 2 to 3 in the energy fitting deserialize method is part of a consistent change across multiple components (ener, dos, dipole, property, and invar_fitting) and is properly supported:
- The version compatibility system is well-implemented with clear error messages
- The model compatibility documentation clearly explains version compatibility and provides migration paths
- The
dp convert-fromtool is available for users to convert models between versions - The changes maintain backward compatibility by accepting version 1 as default (
@version, 1)
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check for documentation of version changes
echo "Checking for documentation of version changes..."
rg -l "version.*2.*3" docs/ || echo "No version change documentation found"
# Check for migration guide
echo "Checking for migration guide..."
rg -l "migrat.*2.*3" docs/ || echo "No migration guide found"
# Check for changelog entries
echo "Checking changelog..."
rg -l "version.*3" CHANGELOG.md || echo "No changelog entry found"
Length of output: 673
Script:
#!/bin/bash
# Check for version-related files or documentation
echo "Searching for documentation files..."
fd -e md -e rst
echo -e "\nChecking for version mentions in Python files..."
rg -l "version.*2.*3" --type py
echo -e "\nChecking the specific version check implementation..."
rg -A 5 "check_version_compatibility" --type py
echo -e "\nChecking for any migration or upgrade related files..."
fd -e py -e md -e rst -e txt | rg -i "migrat|upgrad|chang|version"
Length of output: 64948
Script:
#!/bin/bash
# Check version compatibility implementation
echo "Checking version compatibility implementation..."
cat deepmd/utils/version.py
echo -e "\nChecking version compatibility documentation..."
cat doc/troubleshooting/md-version-compatibility.md
echo -e "\nChecking for version-related changes in the PR file..."
cat deepmd/pd/model/task/ener.py
Length of output: 4135
Script:
#!/bin/bash
# Check git history for version-related changes
echo "Checking recent version-related changes..."
git log -p --since="1 month ago" | grep -A 5 -B 5 "version.*2.*3"
echo -e "\nChecking for model version documentation..."
fd -t f -e md -e rst | xargs grep -l "model.*version"
echo -e "\nChecking for any version-related tests..."
fd -t f -e py | xargs grep -l "test.*version"
Length of output: 527
Script:
#!/bin/bash
# Check for version compatibility tests
echo "Checking version compatibility tests..."
rg -l "test.*version.*compatibility" --type py
echo -e "\nChecking for recent changes in version handling..."
rg -A 5 "check_version_compatibility.*3" --type py
echo -e "\nChecking model documentation for version info..."
cat doc/troubleshooting/model-compatability.md || echo "File not found"
Length of output: 9211
deepmd/pd/model/descriptor/se_atten.py (2)
965-965: Ensure tensor chunking splits correctly
Verify that self.in_proj(query).chunk(3, axis=-1) correctly splits the tensor into three parts along the last axis.
Ensure that embed_dim is divisible by 3 to prevent runtime errors during tensor splitting.
993-1005: Confirm attention weights handling when smoothing is disabled
When self.smooth is False, masking with float('-inf') is used before softmax. Confirm that this approach correctly masks out the padded values without causing numerical issues.
Run tests to ensure that the attention mechanism functions as expected when smoothing is disabled.
deepmd/pd/model/task/fitting.py (2)
57-72: Ensure parameters are correctly shared in share_params method
In the share_params method, when shared_level == 0, the code shares all sub-layers between self and base_class. However, buffers like bias_atom_e and case_embd need to be manually linked, as mentioned in the comment. Ensure that these buffers are correctly shared to avoid inconsistencies during training.
Please verify that bias_atom_e and case_embd are properly shared or intentionally kept independent. If they need to be shared, consider adding code to link these buffers explicitly.
210-219: Ensure case_embd is properly initialized
When dim_case_embd > 0, case_embd is initialized as a zero tensor. Ensure that this is intentional and that case_embd will be correctly set before it is used in the forward pass. Otherwise, it might lead to unexpected behavior.
Please confirm that set_case_embd will be called to set case_embd appropriately before it's used. If not, consider initializing it with meaningful values or raising an error if it's used before being set.
deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
141-147: Addition of set_case_embd method is appropriate
The set_case_embd method correctly adds the functionality to set the case embedding index and aligns with similar implementations in other atomic models.
deepmd/pd/model/model/ener_model.py (1)
44-45:
Incorrect use of squeeze; the result of squeeze is not assigned
The squeeze method does not modify the tensor in place. You need to assign the result back to output_def["force"] to ensure the shape change takes effect. Otherwise, output_def["force"] retains its original shape, which may lead to unexpected behavior downstream.
Apply this diff to fix the issue:
output_def["force"] = out_def_data["energy_derv_r"]
-output_def["force"].squeeze(-2)
+output_def["force"] = output_def["force"].squeeze(-2)⛔ Skipped due to learnings
Learnt from: 1azyking
PR: deepmodeling/deepmd-kit#4169
File: deepmd/pt/model/model/ener_hess_model.py:87-99
Timestamp: 2024-11-12T05:47:21.643Z
Learning: In the deepmd-kit codebase, it's acceptable to call `squeeze` on tensors without assigning the result back, as it matches the existing code patterns (e.g., in `ener_model.py`) and functions correctly.
deepmd/dpmodel/fitting/general_fitting.py (3)
176-179: Initialization of case_embd is appropriate
The initialization of case_embd based on dim_case_embd is correctly implemented.
272-273: Addition of case_embd handling in __setitem__ and __getitem__ methods
The implementation of case_embd in the magic methods __setitem__ and __getitem__ is appropriate and ensures correct assignment and retrieval.
Also applies to: 290-291
446-458: Concatenation of case_embd to the input descriptor
The concatenation of case_embd to the input descriptor xx is correctly handled when dim_case_embd > 0.
deepmd/pd/model/network/network.py (4)
81-81: Ensure atype is a Paddle tensor with a valid place attribute
In the forward method of TypeEmbedNet, atype.place is used to determine the device. Verify that atype is always a Paddle tensor; otherwise, accessing atype.place may raise an AttributeError if atype is, for example, a NumPy array or list.
69-69: Reconsider the commented-out weight initialization code
There is a commented-out line intended for initializing embedding weights in TypeEmbedNet.__init__. If initializing the embedding weights with a normal distribution is necessary, consider uncommenting this line. If it's obsolete, it can be removed to clean up the code.
251-281: Verify matrix dimensions to prevent mismatches during operations
In the change_type_map method, operations involving first_layer_matrix, eye_vector, and concatenations are performed. Ensure that the dimensions align correctly to prevent runtime errors due to dimension mismatches.
242-250:
Verify activation_function is 'Linear' when do_resnet is True
In the change_type_map method of TypeEmbedNetConsistent, there is an assertion that when do_resnet is True, activation_function must be 'Linear'. Ensure this condition is met to prevent assertion errors during runtime.
deepmd/pd/loss/ener.py (2)
98-98: Check consistency of self.has_gf initialization logic
Unlike other self.has_* flags, self.has_gf does not include or inference in its initialization:
self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0All other self.has_* flags include or inference. Is it intended that self.has_gf excludes or inference? If not, the logic should be updated for consistency, ensuring self.has_gf behaves as expected during inference.
158-158: Verify that natoms is always greater than zero to prevent division by zero
The variable atom_norm is computed as 1.0 / natoms. If natoms is zero, this will result in a division by zero error. Ensure that natoms is always greater than zero or add a check to handle the case when natoms is zero.
deepmd/pt/loss/__init__.py (1)
25-31: Approval of updates to the __all__ list
The updates to the __all__ list correctly include the new DOSLoss and reorder the export list to reflect the intended public interface. The changes ensure that all necessary loss classes are publicly accessible.
deepmd/pd/train/training.py (1)
415-415: Verify potential unreachable code
There is a possibility that the code at line 415 is unreachable. Please verify the logic to ensure that this code can be executed as intended.
🧰 Tools
🪛 GitHub Check: CodeQL
[warning] 415-415: Unreachable code
This statement is unreachable.
| @@ -141,7 +145,7 @@ def serialize(self) -> dict: | |||
| @classmethod | |||
| def deserialize(cls, data: dict) -> "GeneralFitting": | |||
| data = data.copy() | |||
| check_version_compatibility(data.pop("@version", 1), 2, 1) | |||
| check_version_compatibility(data.pop("@version", 1), 3, 1) | |||
There was a problem hiding this comment.
Version compatibility change requires documentation.
The version compatibility check has been updated from 2 to 3. This change could affect backward compatibility.
Please ensure that:
- This breaking change is documented in the migration guide
- The version bump follows semantic versioning principles
- Users are provided with clear upgrade instructions
| @@ -73,7 +75,7 @@ def __init__( | |||
| @classmethod | |||
| def deserialize(cls, data: dict) -> "GeneralFitting": | |||
| data = data.copy() | |||
| check_version_compatibility(data.pop("@version", 1), 2, 1) | |||
| check_version_compatibility(data.pop("@version", 1), 3, 1) | |||
There was a problem hiding this comment.
💡 Codebase verification
Version compatibility check reveals inconsistency in polarizability fitting
The version compatibility check reveals that while most fitting classes (EnerFitting, DOSFitting, DipoleFitting, InvarFitting, PropertyFitting) use version 3, the PolarizabilityFitting class uses version 4. This inconsistency needs to be addressed.
deepmd/dpmodel/fitting/polarizability_fitting.py: Uses version 4 while other fitting classes use version 3- Affected classes using version 3:
deepmd/dpmodel/fitting/ener_fitting.pydeepmd/dpmodel/fitting/dos_fitting.pydeepmd/dpmodel/fitting/dipole_fitting.pydeepmd/dpmodel/fitting/invar_fitting.pydeepmd/dpmodel/fitting/property_fitting.py
🔗 Analysis chain
Verify version compatibility across related classes
The version bump from 2 to 3 is appropriate given the addition of case embedding support. However, we should ensure this version change is consistent across all related fitting classes mentioned in the AI summary (DOSFittingNet, PolarFittingNet, etc.).
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check version compatibility in related fitting classes
# Expected: All fitting classes should use version 3 in their deserialize methods
# Search for version checks in fitting classes
rg -A 1 "check_version_compatibility" "deepmd/dpmodel/fitting/"
Length of output: 1940
| @@ -159,7 +161,7 @@ def serialize(self) -> dict: | |||
| @classmethod | |||
| def deserialize(cls, data: dict) -> "GeneralFitting": | |||
| data = data.copy() | |||
| check_version_compatibility(data.pop("@version", 1), 2, 1) | |||
| check_version_compatibility(data.pop("@version", 1), 3, 1) | |||
There was a problem hiding this comment.
💡 Codebase verification
Version compatibility check needs to be updated in polarizability_fitting.py
The version compatibility check is inconsistent across fitting classes:
polarizability_fitting.pyuses version 4- All other fitting classes (
property_fitting.py,ener_fitting.py,dos_fitting.py,dipole_fitting.py,invar_fitting.py) use version 3
This inconsistency should be addressed to maintain uniform versioning across the codebase.
🔗 Analysis chain
Verify version compatibility across fitting classes.
The version check has been updated from 2 to 3, which indicates a breaking change in serialization. Please ensure:
- This change is documented in the migration guide
- Version numbers are consistent across all fitting classes
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Check version compatibility checks in other fitting classes
echo "Checking version compatibility in fitting classes..."
rg -B 1 -A 1 "check_version_compatibility.*version.*\d+.*\d+" "deepmd/dpmodel/fitting/"
Length of output: 1765
| update_residual_init: str = "norm", | ||
| set_davg_zero: bool = True, | ||
| smooth: bool = True, | ||
| exclude_types: list[tuple[int, int]] = [], |
There was a problem hiding this comment.
Replace mutable default arguments with None
Using mutable default arguments (empty list) can lead to unexpected behavior in Python. The same list object will be shared across all instances where the default value is used.
Apply this fix:
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: Optional[list[tuple[int, int]]] = None,And initialize the parameter in the method:
def reinit_exclude(
self,
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: Optional[list[tuple[int, int]]] = None,
) -> None:
+ if exclude_types is None:
+ exclude_types = []
self.exclude_types = exclude_typesAlso applies to: 360-360
🧰 Tools
🪛 Ruff (0.8.0)
81-81: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
| if paddle.in_dynamic_mode(): | ||
| assert isinstance(atype_embd, paddle.Tensor) # for jit | ||
| g1 = self.act(atype_embd) | ||
| ng1 = g1.shape[-1] |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove unused variables
The variables ng1 and idx are assigned but never used in the code.
Apply these changes:
- ng1 = g1.shape[-1]
- for idx, ll in enumerate(self.layers):
+ for _, ll in enumerate(self.layers):Also applies to: 435-435
🧰 Tools
🪛 Ruff (0.8.0)
412-412: Local variable ng1 is assigned to but never used
Remove assignment to unused variable ng1
(F841)
deepmd/pd/loss/ener.py
Outdated
| atom_ener = model_pred["atom_energy"] | ||
| atom_ener_label = label["atom_ener"] | ||
| find_atom_ener = label.get("find_atom_ener", 0.0) |
There was a problem hiding this comment.
Inconsistent key naming: 'atom_energy' vs 'atom_ener'
There is a potential inconsistency in the keys used for atom_energy in model_pred and label. In the model prediction, the key is 'atom_energy', whereas in the label, the key is 'atom_ener'. Similarly, in the label_requirement, the label name is 'atom_ener'. This mismatch may lead to runtime errors when accessing the label data.
Consider standardizing the key names to ensure consistency. For example, change 'atom_ener' to 'atom_energy' in the labels and DataRequirementItem.
Apply the following changes:
At lines 325-327:
-if self.has_ae and "atom_energy" in model_pred and "atom_ener" in label:
- atom_ener = model_pred["atom_energy"]
- atom_ener_label = label["atom_ener"]
+if self.has_ae and "atom_energy" in model_pred and "atom_energy" in label:
+ atom_ener = model_pred["atom_energy"]
+ atom_ener_label = label["atom_energy"]At lines 385-391:
DataRequirementItem(
- "atom_ener",
+ "atom_energy",
ndof=1,
atomic=True,
must=False,
high_prec=False,
)Also applies to: 384-391
deepmd/pd/loss/ener.py
Outdated
| pref_pf = self.limit_pref_pf + (self.start_pref_pf - self.limit_pref_pf) * coef | ||
| pref_gf = self.limit_pref_gf + (self.start_pref_gf - self.limit_pref_gf) * coef | ||
|
|
||
| loss = paddle.zeros([1], dtype=env.GLOBAL_PD_FLOAT_PRECISION).to(env.DEVICE)[0] |
There was a problem hiding this comment.
Invalid method .to() used on PaddlePaddle tensor
The .to() method is not valid for PaddlePaddle tensors. To move the tensor to a specific device in PaddlePaddle, you should specify the device when creating the tensor or use the appropriate method. The current usage will result in an AttributeError.
Apply the following change to correct the issue:
-loss = paddle.zeros([1], dtype=env.GLOBAL_PD_FLOAT_PRECISION).to(env.DEVICE)[0]
+loss = paddle.zeros([1], dtype=env.GLOBAL_PD_FLOAT_PRECISION, place=env.DEVICE)[0]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| loss = paddle.zeros([1], dtype=env.GLOBAL_PD_FLOAT_PRECISION).to(env.DEVICE)[0] | |
| loss = paddle.zeros([1], dtype=env.GLOBAL_PD_FLOAT_PRECISION, place=env.DEVICE)[0] |
deepmd/pd/loss/ener.py
Outdated
| repeat=3, | ||
| ) | ||
| ) | ||
| if self.has_gf > 0: |
There was a problem hiding this comment.
Logical error: comparing boolean self.has_gf with integer
In line 403, self.has_gf is a boolean value, but it is being compared to an integer using > 0. This comparison may not behave as intended because True > 0 evaluates to True, but explicitly comparing a boolean to an integer is discouraged.
Replace the condition with the appropriate boolean check.
Apply the following change:
-if self.has_gf > 0:
+if self.has_gf:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if self.has_gf > 0: | |
| if self.has_gf: |
deepmd/pd/train/training.py
Outdated
| ) # line buffered | ||
| if SAMPLER_RECORD: | ||
| record_file = f"Sample_rank_{self.rank}.txt" | ||
| fout1 = open(record_file, mode="w", buffering=1) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use a context manager to ensure the file is properly closed
Currently, the file fout1 may not be properly closed if an exception occurs before reaching fout1.close(). Using a context manager (with statement) ensures that the file is closed correctly.
Refactor the code to use a with statement:
if SAMPLER_RECORD:
record_file = f"Sample_rank_{self.rank}.txt"
- fout1 = open(record_file, mode="w", buffering=1)
+ with open(record_file, mode="w", buffering=1) as fout1:
+ # Rest of the code that uses fout1
+ ...You will need to adjust the code to ensure that all usages of fout1 are within the context manager.
Committable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 Ruff (0.8.0)
646-646: Use a context manager for opening files
(SIM115)
deepmd/pd/train/training.py
Outdated
| open( | ||
| self.disp_file, | ||
| mode="w" if not self.restart_training else "a", | ||
| buffering=1, | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Use a context manager to ensure the file is properly closed
Currently, the file fout may not be properly closed if an exception occurs before reaching fout.close(). Using a context manager (with statement) ensures that the file is closed correctly.
Refactor the code to use a with statement:
- fout = (
- open(
- self.disp_file,
- mode="w" if not self.restart_training else "a",
- buffering=1,
- )
- if self.rank == 0
- else None
- ) # line buffered
+ if self.rank == 0:
+ with open(
+ self.disp_file,
+ mode="w" if not self.restart_training else "a",
+ buffering=1,
+ ) as fout:
+ # Rest of the code that uses fout
+ ...You will need to adjust the code to ensure that all usages of fout are within the context manager.
Committable suggestion skipped: line range outside the PR's diff.
🧰 Tools
🪛 Ruff (0.8.0)
636-636: Use a context manager for opening files
(SIM115)
HydrogenSulfate
force-pushed
the
add_dpa2_v2
branch
from
64fc62e to
79ed0f0
Compare
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Outside diff range and nitpick comments (6)
deepmd/pd/model/descriptor/se_t_tebd.py (3)
823-827: Consider using einsum for better readabilityThe commented-out einsum operations have been replaced with manual tensor operations. While this might be for performance reasons, consider documenting why einsum was avoided or if it's still a valid alternative.
Original einsum operations:
# env_ij = paddle.einsum("ijm,ikm->ijk", rr_i, rr_j) # res_ij = paddle.einsum("ijk,ijkm->im", env_ij, gg)Also applies to: 909-913
386-386: Remove unused variable assignmentsSeveral variables are assigned but never used in the code:
env_matin the deserialize methodnallin multiple locationsApply these changes:
- env_mat = data.pop("env_mat") + data.pop("env_mat") # Required for compatibility - nall = extended_coord.reshape([nframes, -1]).shape[1] // 3 + _ = extended_coord.reshape([nframes, -1]).shape[1] // 3 # Kept for validation - nall = extended_coord.reshape([nb, -1, 3]).shape[1] + _ = extended_coord.reshape([nb, -1, 3]).shape[1] # Kept for validationAlso applies to: 454-454, 787-787
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 386-386: Unused local variable
Variable env_mat is not used.
831-848: Consider optimizing tensor operationsThe tensor operations in the
tebd_input_mode == "concat"branch involve multiple reshapes and concatenations. Consider:
- Pre-allocating tensors for frequently used shapes
- Combining multiple reshapes into a single operation where possible
- Using in-place operations where applicable
Example optimization:
# Pre-allocate expanded tensors with correct shapes nlist_tebd_i = paddle.empty([nfnl, self.nnei, self.nnei, self.tebd_dim]) nlist_tebd_j = paddle.empty([nfnl, self.nnei, self.nnei, self.tebd_dim]) # Use a single reshape operation nlist_tebd = atype_tebd_nlist.reshape([nfnl, nnei, self.tebd_dim]) # Use in-place expand operations nlist_tebd.expand_([nfnl, self.nnei, self.nnei, -1], out=nlist_tebd_i) nlist_tebd.expand_([nfnl, self.nnei, self.nnei, -1], out=nlist_tebd_j)deepmd/pd/model/descriptor/repformer_layer.py (3)
41-82: Fix docstring parameter descriptionThe docstring for
_scaleparameter inget_residualfunction describes it as "The initial scale of the residual tensor" but doesn't fully explain its usage in different initialization modes.Add mode-specific details to the docstring:
_scale - The initial scale of the residual tensor. See `_mode` for details. + The initial scale of the residual tensor: + - In "norm" mode: Used as standard deviation for normal distribution + - In "const" mode: Used as constant value for all elements
580-616: Improve parameter organization and input validationThe class initialization has too many parameters and uses assertions for validation.
Consider these improvements:
- Use a configuration class to group parameters:
@dataclass class RepformerConfig: rcut: float rcut_smth: float sel: int ntypes: int g1_dim: int = 128 g2_dim: int = 16 axis_neuron: int = 4 # ... other parameters with defaults
- Replace assertions with proper validation:
- assert len(sel) == 1 + if len(sel) != 1: + raise ValueError("sel must be a list with exactly one element")
1413-1484: Enhance error handling in deserializationThe deserialize method could benefit from more robust error handling for malformed data.
Consider adding validation for required fields and type checking:
@classmethod def deserialize(cls, data: dict) -> "RepformerLayer": data = data.copy() + # Validate required fields + required_fields = ["@version", "@class", "linear1"] + missing_fields = [f for f in required_fields if f not in data] + if missing_fields: + raise ValueError(f"Missing required fields: {missing_fields}") + check_version_compatibility(data.pop("@version"), 2, 1)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (3)
deepmd/pd/model/descriptor/__init__.py(1 hunks)deepmd/pd/model/descriptor/repformer_layer.py(1 hunks)deepmd/pd/model/descriptor/se_t_tebd.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pd/model/descriptor/init.py
🧰 Additional context used
🪛 Ruff (0.8.0)
deepmd/pd/model/descriptor/se_t_tebd.py
129-129: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
136-136: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
387-387: Local variable env_mat is assigned to but never used
Remove assignment to unused variable env_mat
(F841)
454-454: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
521-521: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
528-528: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
709-713: Use ternary operator sampled = merged() if callable(merged) else merged instead of if-else-block
(SIM108)
733-733: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
787-787: Local variable nall is assigned to but never used
Remove assignment to unused variable nall
(F841)
deepmd/pd/model/descriptor/repformer_layer.py
967-967: Local variable ng2 is assigned to but never used
Remove assignment to unused variable ng2
(F841)
1160-1163: Use ternary operator gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None instead of if-else-block
Replace if-else-block with gg1 = _make_nei_g1(g1_ext, nlist) if cal_gg1 else None
(SIM108)
1281-1281: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
🪛 GitHub Check: CodeQL
deepmd/pd/model/descriptor/se_t_tebd.py
[notice] 386-386: Unused local variable
Variable env_mat is not used.
[notice] 451-451: Unused local variable
Variable nall is not used.
[notice] 770-770: Unused local variable
Variable nall is not used.
deepmd/pd/model/descriptor/repformer_layer.py
[notice] 967-967: Unused local variable
Variable ng2 is not used.
[notice] 1277-1277: Unused local variable
Variable nitem is not used.
🔇 Additional comments (2)
deepmd/pd/model/descriptor/se_t_tebd.py (1)
129-129:
Fix mutable default arguments in constructors
Using mutable objects (lists) as default arguments can lead to unexpected behavior because the same list object is shared across all function calls.
Apply these changes to all affected constructors:
# In DescrptSeTTebd.__init__
- neuron: list = [2, 4, 8],
+ neuron: Optional[list] = None,
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: Optional[list[tuple[int, int]]] = None,
# In DescrptBlockSeTTebd.__init__
- neuron: list = [25, 50, 100],
+ neuron: Optional[list] = None,
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: Optional[list[tuple[int, int]]] = None,
# In reinit_exclude
- exclude_types: list[tuple[int, int]] = [],
+ exclude_types: Optional[list[tuple[int, int]]] = None,Add initialization in the constructors:
def __init__(self, ...):
if neuron is None:
neuron = [2, 4, 8] # or [25, 50, 100] for DescrptBlockSeTTebd
if exclude_types is None:
exclude_types = []Also applies to: 136-136, 521-521, 528-528, 733-733
🧰 Tools
🪛 Ruff (0.8.0)
129-129: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pd/model/descriptor/repformer_layer.py (1)
155-243: 🛠️ Refactor suggestion
Add input validation and remove unused variable
The forward method could benefit from explicit input validation and has an unused variable.
Consider these improvements:
- Remove unused variable:
- nb, nloc, nnei, _ = g2.shape
- ng2 = g2.shape[-1]
+ nb, nloc, nnei, _ = g2.shape- Add input shape validation:
+ # Validate input shapes
+ if g2.ndim != 4 or h2.ndim != 4:
+ raise ValueError("g2 and h2 must be 4-dimensional tensors")
+ if h2.shape[-1] != 3:
+ raise ValueError("h2 last dimension must be 3")
+ if g2.shape[:-1] != h2.shape[:-1]:
+ raise ValueError("g2 and h2 must have compatible shapes")Likely invalid or redundant comment.
| def list_update_res_residual( | ||
| self, update_list: list[paddle.Tensor], update_name: str = "g1" | ||
| ) -> paddle.Tensor: | ||
| nitem = len(update_list) | ||
| uu = update_list[0] | ||
| # make jit happy | ||
| if update_name == "g1": | ||
| for ii, vv in enumerate(self.g1_residual): | ||
| uu = uu + vv * update_list[ii + 1] | ||
| elif update_name == "g2": | ||
| for ii, vv in enumerate(self.g2_residual): | ||
| uu = uu + vv * update_list[ii + 1] | ||
| elif update_name == "h2": | ||
| for ii, vv in enumerate(self.h2_residual): | ||
| uu = uu + vv * update_list[ii + 1] | ||
| else: | ||
| raise NotImplementedError | ||
| return uu | ||
|
|
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove unused variable
The list_update_res_residual method has an unused variable.
Remove the unused variable:
def list_update_res_residual(
self, update_list: list[paddle.Tensor], update_name: str = "g1"
) -> paddle.Tensor:
- nitem = len(update_list)
uu = update_list[0]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def list_update_res_residual( | |
| self, update_list: list[paddle.Tensor], update_name: str = "g1" | |
| ) -> paddle.Tensor: | |
| nitem = len(update_list) | |
| uu = update_list[0] | |
| # make jit happy | |
| if update_name == "g1": | |
| for ii, vv in enumerate(self.g1_residual): | |
| uu = uu + vv * update_list[ii + 1] | |
| elif update_name == "g2": | |
| for ii, vv in enumerate(self.g2_residual): | |
| uu = uu + vv * update_list[ii + 1] | |
| elif update_name == "h2": | |
| for ii, vv in enumerate(self.h2_residual): | |
| uu = uu + vv * update_list[ii + 1] | |
| else: | |
| raise NotImplementedError | |
| return uu | |
| def list_update_res_residual( | |
| self, update_list: list[paddle.Tensor], update_name: str = "g1" | |
| ) -> paddle.Tensor: | |
| uu = update_list[0] | |
| # make jit happy | |
| if update_name == "g1": | |
| for ii, vv in enumerate(self.g1_residual): | |
| uu = uu + vv * update_list[ii + 1] | |
| elif update_name == "g2": | |
| for ii, vv in enumerate(self.g2_residual): | |
| uu = uu + vv * update_list[ii + 1] | |
| elif update_name == "h2": | |
| for ii, vv in enumerate(self.h2_residual): | |
| uu = uu + vv * update_list[ii + 1] | |
| else: | |
| raise NotImplementedError | |
| return uu |
🧰 Tools
🪛 Ruff (0.8.0)
1281-1281: Local variable nitem is assigned to but never used
Remove assignment to unused variable nitem
(F841)
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4418 +/- ##
==========================================
+ Coverage 84.39% 84.58% +0.19%
==========================================
Files 670 675 +5
Lines 62130 63496 +1366
Branches 3487 3486 -1
==========================================
+ Hits 52433 53709 +1276
- Misses 8571 8661 +90
Partials 1126 1126 ☔ View full report in Codecov by Sentry. |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Outside diff range and nitpick comments (2)
source/tests/pd/model/test_descriptor_dpa2.py (2)
24-118: Add docstring and improve test data organization.While the test setup is thorough, consider these improvements for better maintainability:
- Add a docstring explaining the test setup and the meaning of the test data
- Consider moving the hardcoded test values to separate data files
- Add comments explaining the significance of the reference values
Example docstring:
class TestDPA2(unittest.TestCase): def setUp(self): + """Initialize test data for DPA2 descriptor testing. + + Sets up: + - Simulation cell parameters + - Atomic coordinates + - Atom types + - Reference descriptor values + - Paths to model parameters and embeddings + """
119-186: Enhance test coverage and improve test organization.The test covers the happy path but could be improved:
- Split into separate test methods for different configurations
- Add error cases and edge cases
- Add input validation tests
Consider refactoring like this:
def test_descriptor_basic_config(self): """Test DPA2 descriptor with basic configuration.""" # Current test logic with concat_output_tebd=False def test_descriptor_with_concat_output(self): """Test DPA2 descriptor with concatenated output.""" # Current test logic with concat_output_tebd=True def test_descriptor_invalid_inputs(self): """Test DPA2 descriptor with invalid inputs.""" # Add tests for invalid coordinates, types, etc. def test_descriptor_edge_cases(self): """Test DPA2 descriptor with edge cases.""" # Add tests for empty system, single atom, etc.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
📒 Files selected for processing (1)
source/tests/pd/model/test_descriptor_dpa2.py(1 hunks)
🔇 Additional comments (1)
source/tests/pd/model/test_descriptor_dpa2.py (1)
1-23: LGTM! Well-organized imports following best practices.
The imports are properly organized (standard library → third-party → local) and use modern Python practices like Path for file operations.
| def translate_type_embd_dicts_to_dpa2( | ||
| target_dict, | ||
| source_dict, | ||
| type_embd_dict, | ||
| ): | ||
| all_keys = list(target_dict.keys()) | ||
| record = [False for ii in all_keys] | ||
| for kk, vv in source_dict.items(): | ||
| record[all_keys.index(kk)] = True | ||
| target_dict[kk] = vv | ||
| assert len(type_embd_dict.keys()) == 2 | ||
| it = iter(type_embd_dict.keys()) | ||
| for _ in range(2): | ||
| kk = next(it) | ||
| tk = "type_embedding." + kk | ||
| record[all_keys.index(tk)] = True | ||
| target_dict[tk] = type_embd_dict[kk] | ||
| record[all_keys.index("repinit.compress_data.0")] = True | ||
| record[all_keys.index("repinit.compress_info.0")] = True | ||
| assert all(record) | ||
| return target_dict |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add type hints and improve documentation.
The utility function needs better documentation and type safety:
- Add type hints for parameters and return value
- Add docstring explaining the function's purpose
- Improve assertion message
Apply these improvements:
def translate_type_embd_dicts_to_dpa2(
- target_dict,
- source_dict,
- type_embd_dict,
+ target_dict: dict,
+ source_dict: dict,
+ type_embd_dict: dict,
+) -> dict:
+ """Translate type embedding dictionaries for DPA2 descriptor.
+
+ Args:
+ target_dict: Target state dictionary to update
+ source_dict: Source state dictionary containing model parameters
+ type_embd_dict: Dictionary containing type embeddings
+
+ Returns:
+ Updated target dictionary with translated embeddings
+
+ Raises:
+ AssertionError: If type_embd_dict doesn't contain exactly 2 keys
+ AssertionError: If any required keys are missing in the final dictionary
+ """
all_keys = list(target_dict.keys())
record = [False for ii in all_keys]
for kk, vv in source_dict.items():
record[all_keys.index(kk)] = True
target_dict[kk] = vv
- assert len(type_embd_dict.keys()) == 2
+ assert len(type_embd_dict.keys()) == 2, "Expected exactly 2 type embeddings"
it = iter(type_embd_dict.keys())
for _ in range(2):
kk = next(it)
tk = "type_embedding." + kk
record[all_keys.index(tk)] = True
target_dict[tk] = type_embd_dict[kk]
record[all_keys.index("repinit.compress_data.0")] = True
record[all_keys.index("repinit.compress_info.0")] = True
- assert all(record)
+ assert all(record), "Missing required keys in the final dictionary"
return target_dict📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def translate_type_embd_dicts_to_dpa2( | |
| target_dict, | |
| source_dict, | |
| type_embd_dict, | |
| ): | |
| all_keys = list(target_dict.keys()) | |
| record = [False for ii in all_keys] | |
| for kk, vv in source_dict.items(): | |
| record[all_keys.index(kk)] = True | |
| target_dict[kk] = vv | |
| assert len(type_embd_dict.keys()) == 2 | |
| it = iter(type_embd_dict.keys()) | |
| for _ in range(2): | |
| kk = next(it) | |
| tk = "type_embedding." + kk | |
| record[all_keys.index(tk)] = True | |
| target_dict[tk] = type_embd_dict[kk] | |
| record[all_keys.index("repinit.compress_data.0")] = True | |
| record[all_keys.index("repinit.compress_info.0")] = True | |
| assert all(record) | |
| return target_dict | |
| def translate_type_embd_dicts_to_dpa2( | |
| target_dict: dict, | |
| source_dict: dict, | |
| type_embd_dict: dict, | |
| ) -> dict: | |
| """Translate type embedding dictionaries for DPA2 descriptor. | |
| Args: | |
| target_dict: Target state dictionary to update | |
| source_dict: Source state dictionary containing model parameters | |
| type_embd_dict: Dictionary containing type embeddings | |
| Returns: | |
| Updated target dictionary with translated embeddings | |
| Raises: | |
| AssertionError: If type_embd_dict doesn't contain exactly 2 keys | |
| AssertionError: If any required keys are missing in the final dictionary | |
| """ | |
| all_keys = list(target_dict.keys()) | |
| record = [False for ii in all_keys] | |
| for kk, vv in source_dict.items(): | |
| record[all_keys.index(kk)] = True | |
| target_dict[kk] = vv | |
| assert len(type_embd_dict.keys()) == 2, "Expected exactly 2 type embeddings" | |
| it = iter(type_embd_dict.keys()) | |
| for _ in range(2): | |
| kk = next(it) | |
| tk = "type_embedding." + kk | |
| record[all_keys.index(tk)] = True | |
| target_dict[tk] = type_embd_dict[kk] | |
| record[all_keys.index("repinit.compress_data.0")] = True | |
| record[all_keys.index("repinit.compress_info.0")] = True | |
| assert all(record), "Missing required keys in the final dictionary" | |
| return target_dict |
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
source/tests/pd/test_multitask.py(2 hunks)source/tests/pd/test_training.py(2 hunks)
🔇 Additional comments (3)
source/tests/pd/test_training.py (2)
27-27: LGTM!
The import statement follows the established pattern and is properly placed with other model imports.
199-213: LGTM! Verify test data availability.
The test class implementation follows the established pattern and maintains consistency with other test classes. The implementation looks correct and well-structured.
Let's verify the test data availability:
#!/bin/bash
# Description: Verify that the required test data exists
# Test data path: water/data/data_0
# Check if the test data directory exists
if [ -d "source/tests/pd/water/data/data_0" ]; then
echo "Test data directory exists"
# List contents to verify structure
ls -la source/tests/pd/water/data/data_0
else
echo "Warning: Test data directory not found"
fi
# Check if the configuration file exists
if [ -f "source/tests/pd/water/se_atten.json" ]; then
echo "Configuration file exists"
else
echo "Warning: Configuration file not found"
fisource/tests/pd/test_multitask.py (1)
33-33: LGTM!
The import statement follows the established pattern and is properly placed with other model imports.
source/tests/pd/test_multitask.py
Outdated
| class TestMultiTaskDPA2(unittest.TestCase, MultiTaskTrainTest): | ||
| def setUp(self) -> None: | ||
| multitask_DPA2 = deepcopy(multitask_template) | ||
| multitask_DPA2["model"]["shared_dict"]["my_descriptor"] = model_dpa2[ | ||
| "descriptor" | ||
| ] | ||
| data_file = [str(Path(__file__).parent / "water/data/data_0")] | ||
| self.stat_files = "DPA2" | ||
| os.makedirs(self.stat_files, exist_ok=True) | ||
| self.config = multitask_DPA2 | ||
|
|
There was a problem hiding this comment.
Complete the test class implementation.
The test class is incomplete compared to other similar test classes (TestMultiTaskSeA, TestMultiTaskDPA1). Several required configurations and methods are missing:
- Configuration for model_2 is missing
- Training and validation settings are not complete
- Missing origin_config and shared_links setup
- Missing tearDown method
Here's how to complete the implementation:
class TestMultiTaskDPA2(unittest.TestCase, MultiTaskTrainTest):
def setUp(self) -> None:
multitask_DPA2 = deepcopy(multitask_template)
multitask_DPA2["model"]["shared_dict"]["my_descriptor"] = model_dpa2[
"descriptor"
]
data_file = [str(Path(__file__).parent / "water/data/data_0")]
self.stat_files = "DPA2"
os.makedirs(self.stat_files, exist_ok=True)
self.config = multitask_DPA2
+ self.config["training"]["data_dict"]["model_1"]["training_data"]["systems"] = (
+ data_file
+ )
+ self.config["training"]["data_dict"]["model_1"]["validation_data"][
+ "systems"
+ ] = data_file
+ self.config["training"]["data_dict"]["model_1"]["stat_file"] = (
+ f"{self.stat_files}/model_1"
+ )
+ self.config["training"]["data_dict"]["model_2"]["training_data"]["systems"] = (
+ data_file
+ )
+ self.config["training"]["data_dict"]["model_2"]["validation_data"][
+ "systems"
+ ] = data_file
+ self.config["training"]["data_dict"]["model_2"]["stat_file"] = (
+ f"{self.stat_files}/model_2"
+ )
+ self.config["training"]["numb_steps"] = 1
+ self.config["training"]["save_freq"] = 1
+ self.origin_config = deepcopy(self.config)
+ self.config["model"], self.shared_links = preprocess_shared_params(
+ self.config["model"]
+ )
+
+ def tearDown(self) -> None:
+ MultiTaskTrainTest.tearDown(self)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| class TestMultiTaskDPA2(unittest.TestCase, MultiTaskTrainTest): | |
| def setUp(self) -> None: | |
| multitask_DPA2 = deepcopy(multitask_template) | |
| multitask_DPA2["model"]["shared_dict"]["my_descriptor"] = model_dpa2[ | |
| "descriptor" | |
| ] | |
| data_file = [str(Path(__file__).parent / "water/data/data_0")] | |
| self.stat_files = "DPA2" | |
| os.makedirs(self.stat_files, exist_ok=True) | |
| self.config = multitask_DPA2 | |
| class TestMultiTaskDPA2(unittest.TestCase, MultiTaskTrainTest): | |
| def setUp(self) -> None: | |
| multitask_DPA2 = deepcopy(multitask_template) | |
| multitask_DPA2["model"]["shared_dict"]["my_descriptor"] = model_dpa2[ | |
| "descriptor" | |
| ] | |
| data_file = [str(Path(__file__).parent / "water/data/data_0")] | |
| self.stat_files = "DPA2" | |
| os.makedirs(self.stat_files, exist_ok=True) | |
| self.config = multitask_DPA2 | |
| self.config["training"]["data_dict"]["model_1"]["training_data"]["systems"] = ( | |
| data_file | |
| ) | |
| self.config["training"]["data_dict"]["model_1"]["validation_data"][ | |
| "systems" | |
| ] = data_file | |
| self.config["training"]["data_dict"]["model_1"]["stat_file"] = ( | |
| f"{self.stat_files}/model_1" | |
| ) | |
| self.config["training"]["data_dict"]["model_2"]["training_data"]["systems"] = ( | |
| data_file | |
| ) | |
| self.config["training"]["data_dict"]["model_2"]["validation_data"][ | |
| "systems" | |
| ] = data_file | |
| self.config["training"]["data_dict"]["model_2"]["stat_file"] = ( | |
| f"{self.stat_files}/model_2" | |
| ) | |
| self.config["training"]["numb_steps"] = 1 | |
| self.config["training"]["save_freq"] = 1 | |
| self.origin_config = deepcopy(self.config) | |
| self.config["model"], self.shared_links = preprocess_shared_params( | |
| self.config["model"] | |
| ) | |
| def tearDown(self) -> None: | |
| MultiTaskTrainTest.tearDown(self) |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| repinit_variable.pop("embeddings_strip") | ||
| ) | ||
| obj.repinit["davg"] = t_cvt(statistic_repinit["davg"]) | ||
| obj.repinit["dstd"] = t_cvt(statistic_repinit["dstd"]) |
Check warning
Code scanning / CodeQL
Variable defined multiple times Warning
| NetworkCollection.deserialize( | ||
| repinit_three_body_variable.pop("embeddings_strip") | ||
| ) | ||
| ) |
Check warning
Code scanning / CodeQL
Variable defined multiple times Warning
| obj.repformers.layers = paddle.nn.LayerList( | ||
| [RepformerLayer.deserialize(layer) for layer in repformer_layers] | ||
| ) | ||
| return obj |
Check notice
Code scanning / CodeQL
Unused local variable Note
| # h2: nb x nloc x nnei x 3 | ||
| # msk: nb x nloc x nnei | ||
| nb, nloc, nnei, _ = g2.shape | ||
| ng2 = g2.shape[-1] |
Check notice
Code scanning / CodeQL
Unused local variable Note
| for ii in range(1, nitem): | ||
| uu = uu + scale * update_list[ii] | ||
| return uu | ||
|
|
Check notice
Code scanning / CodeQL
Unused local variable Note
| [False, True], # use_econf_tebd | ||
| [True], # new sub-structures (use_sqrt_nnei, g1_out_conv, g1_out_mlp) | ||
| ): | ||
| dtype = PRECISION_DICT[prec] |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| ) | ||
| dd0.repformers.stddev = paddle.to_tensor(dstd_2, dtype=dtype).to( | ||
| device=env.DEVICE | ||
| ) |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| def _test_unused(self, model_params): | ||
| self.model = get_model(model_params).to(env.DEVICE) | ||
| natoms = 5 | ||
| generator = paddle.seed(GLOBAL_SEED) |
Check notice
Code scanning / CodeQL
Unused local variable Note test
| coord = paddle.rand([natoms, 3], dtype=dtype).to(device=env.DEVICE) | ||
| coord = paddle.matmul(coord, cell) | ||
| atype = paddle.to_tensor([0, 0, 0, 1, 1]).to(env.DEVICE) | ||
| idx_perm = [1, 0, 4, 3, 2] |
Check notice
Code scanning / CodeQL
Unused local variable Note test
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
deepmd/pd/model/task/fitting.py (1)
214-215: Improve tensor initialization and device placement.The change from using
placeargument to.to(device)is good for consistency with PyTorch-style device placement. However, consider usingpaddle.zeros_likefor better efficiency when reinitializing tensors.- paddle.zeros(self.dim_case_embd, dtype=self.prec).to(device=device), + paddle.zeros_like(paddle.empty([self.dim_case_embd], dtype=self.prec)).to(device=device),deepmd/pd/utils/multi_task.py (2)
99-101: LGTM! Consider adding parameter type hints.The addition of the return type annotation
-> Noneis a good practice as it clearly indicates the function modifies the input in-place. Consider also adding type hints for the parameters to further improve code clarity.def replace_one_item( - params_dict, key_type, key_in_dict, suffix="", index=None + params_dict: dict, + key_type: str, + key_in_dict: str, + suffix: str = "", + index: Optional[int] = None ) -> None:
99-101: Document sharing levels in function docstring.The function handles sophisticated parameter sharing with different sharing levels (e.g., "my_des1:1"), but this feature isn't well-documented. Consider adding a docstring explaining:
- The sharing level syntax (e.g., ":1")
- How sharing levels affect parameter sharing
- Special handling of hybrid descriptors
def replace_one_item( params_dict, key_type, key_in_dict, suffix="", index=None ) -> None: """Replace parameters with shared values based on sharing configuration. Args: params_dict: Dictionary to modify with shared parameters key_type: Type of parameter (e.g., "descriptor", "fitting_net") key_in_dict: Key in shared_dict, optionally with sharing level (e.g., "my_des1:1") suffix: Optional suffix for hybrid models index: Optional index for hybrid model list Sharing levels: - No suffix (e.g., "my_des1"): Share all parameters (level 0) - With suffix (e.g., "my_des1:1"): Share subset of parameters (level 1+) """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
deepmd/pd/model/task/fitting.py(1 hunks)deepmd/pd/utils/multi_task.py(1 hunks)source/tests/pd/model/water/multitask.json(1 hunks)source/tests/pd/model/water/multitask_sharefit.json(2 hunks)source/tests/pd/test_multitask.py(4 hunks)
🔇 Additional comments (5)
deepmd/pd/model/task/fitting.py (1)
214-215: Verify case embedding usage patterns across the codebase.
Let's verify how case embeddings are used across the codebase to ensure this change doesn't impact other components.
✅ Verification successful
Based on the search results, I can now generate the final response:
Change from zeros to eye initialization appears safe
The change from paddle.zeros() to commented-out paddle.eye()[0] for case embeddings initialization is safe because:
- The actual case embedding used during inference is overridden in the
get_case_embdmethod (line 273) with:
self.case_embd = paddle.eye(self.dim_case_embd, dtype=self.prec).to(device)[case_idx]-
The initial value is only used as a placeholder during model initialization and doesn't affect the model's behavior since it gets replaced with the proper one-hot encoded case embedding during runtime.
-
The device placement is consistent in both versions since both use the same
.to(device=device)pattern.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Search for case embedding usage patterns
echo "Searching for case embedding usage..."
rg -l "case_embd" --type py
# Search for potential device placement patterns
echo "Checking device placement patterns..."
rg "\.to\(device" --type py -A 2
Length of output: 39817
source/tests/pd/model/water/multitask_sharefit.json (1)
Line range hint 94-118: LGTM! Path updates are consistent with Paddle backend support.
The changes correctly update the data paths from PyTorch (pt) to Paddle (pd) backend.
source/tests/pd/test_multitask.py (2)
324-324: False positive: data_file variable is used.
The static analysis tool incorrectly flags data_file as unused. This variable is used in the configuration setup for both training and validation data systems.
Line range hint 239-394: LGTM! Test classes are well-structured and complete.
The new test classes (TestMultiTaskSeASharefit, TestMultiTaskDPA2, TestMultiTaskDPA2Tebd) are properly implemented with:
- Correct inheritance from unittest.TestCase and MultiTaskTrainTest
- Complete setUp methods with proper configuration initialization
- Proper tearDown methods
- Consistent test patterns with existing test classes
deepmd/pd/utils/multi_task.py (1)
99-101: Verify DPA-2 descriptor registration.
Since this PR adds DPA-2 support, we should verify that the DPA-2 descriptor is properly registered and can be handled by the parameter sharing mechanism.
✅ Verification successful
Let me gather more information to verify the DPA-2 descriptor registration and implementation.
Based on the search results, I can now provide a final response:
DPA-2 descriptor is properly integrated and registered
The codebase shows comprehensive DPA-2 descriptor implementation across multiple backends:
- Paddle implementation in
deepmd/pd/model/descriptor/dpa2.py - PyTorch implementation in
deepmd/pt/model/descriptor/dpa2.py - JAX implementation in
deepmd/jax/descriptor/dpa2.py - Base implementation in
deepmd/dpmodel/descriptor/dpa2.py
The descriptor is also properly used in configuration files (*.json, *.yaml) and has corresponding test cases. The parameter sharing mechanism in multi_task.py will work correctly with DPA-2 descriptors.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Verify DPA-2 descriptor registration and usage in parameter sharing
# Check if DPA-2 descriptor is registered
rg -A 5 "register_descriptor.*dpa2"
# Check for DPA-2 usage in configuration files
rg -g "*.json" -g "*.yaml" "type.*dpa2"
# Verify descriptor class implementation
ast-grep --pattern 'class DescrptDPA2(BaseDescriptor):
$$$
'
Length of output: 807
Script:
#!/bin/bash
# Check for DPA-2 descriptor registration in registry files
rg -A 5 "register.*descriptor.*dpa"
# Look for DPA-2 descriptor class definition
rg -A 10 "class.*Descrpt.*DPA2?"
# Check imports and usage of DPA-2 descriptor
rg "from.*descrpt.*dpa|import.*descrpt.*dpa"
Length of output: 17188
| paddle.zeros(self.dim_case_embd, dtype=self.prec).to(device=device), | ||
| # paddle.eye(self.dim_case_embd, dtype=self.prec).to(device=device)[0], |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Add input validation in set_case_embd method.
The set_case_embd method uses the case_embd buffer initialized on these lines, but lacks proper input validation which could lead to index out of bounds errors.
def set_case_embd(self, case_idx: int):
+ if not isinstance(case_idx, int):
+ raise TypeError(f"case_idx must be an integer, got {type(case_idx)}")
+ if case_idx < 0 or case_idx >= self.dim_case_embd:
+ raise ValueError(
+ f"case_idx must be between 0 and {self.dim_case_embd-1}, got {case_idx}"
+ )
self.case_embd = paddle.eye(self.dim_case_embd, dtype=self.prec).to(device)[
case_idx
]Also applies to: 392-394
| 92, | ||
| 4 |
There was a problem hiding this comment.
💡 Codebase verification
Based on the search results, I can generate the final response since I have enough information to verify the DPA-2 selection criteria change:
The selection criteria change from [46, 92] to [92, 4] appears inconsistent with other DPA-2 implementations
Looking at the codebase:
- The original [46, 92] selection criteria is consistently used across water model configurations in multiple files (e.g. examples/water/se_e2_a/input.json, source/tests/pd/model/water/se_e2_a.json)
- The new [92, 4] selection appears only in multitask-related files and seems to deviate from the established pattern
- This change could potentially break compatibility with the DPA-2 descriptor implementation since it significantly alters the network architecture
🔗 Analysis chain
Verify the updated selection criteria for DPA-2.
The sel array has been modified from [46, 92] to [92, 4]. Please ensure this change is compatible with the DPA-2 descriptor implementation.
🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash
# Description: Check if the new selection criteria is used consistently across DPA-2 related files
# Search for sel array definitions in descriptor configurations
rg -A 5 '"sel":\s*\['
# Search for any hardcoded selection criteria in the DPA-2 implementation
ast-grep --pattern 'sel = [$$]'
Length of output: 41995
|
@njzjz This PR is ready for review |
Support DPA-2 in paddle backend. This PR will be updated after #4414 is merged.
Training curve:
Accuracy test(left: paddle, right: torch):
Ralated optimization of Paddle framework:
paddle.whereandpaddle.where_in eager mode PaddlePaddle/Paddle#69556Summary by CodeRabbit
New Features
DescrptDPA2,DescrptBlockRepformers,DescrptSeTTebd, andDescrptBlockSeTTebd.Bug Fixes
Documentation
Tests
DescrptDPA2and multitask training scenarios.