pd(feat): support python inference with DP class #4987
Conversation
…to pd_deepeval
There was a problem hiding this comment.
Pull Request Overview
This PR introduces comprehensive support for Python inference with the DP class, primarily focusing on enhancing the Paddle backend implementation. The changes enable loading and evaluating static models from .json files, expand model type handling to support various output types, and improve output shape determination and evaluation logic.
- Adds static model inference support with Paddle inference engine integration for
.jsonfiles - Expands model type detection to support DOS, dipole, polar, global polar, WFC, and property models
- Enhances output variable handling and implements previously unimplemented evaluation methods
- Updates model freezing to support additional methods and atomic virial computation
Reviewed Changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| source/tests/pd/model/test_deeppot.py | New comprehensive test suite for DeepPot functionality including frozen model testing |
| source/api_cc/src/DeepPotPD.cc | C++ API improvements for tensor handling and parameter naming consistency |
| doc/third-party/ase.md | Documentation updates adding Paddle backend examples alongside existing TensorFlow examples |
| Multiple descriptor files | Buffer registration for static model compatibility across descriptor implementations |
| deepmd/pd/infer/deep_eval.py | Core implementation of static model support and expanded inference capabilities |
| deepmd/pd/entrypoints/main.py | Enhanced model freezing with additional method support and atomic virial parameter |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughWalkthroughAdds persistent buffer tensors and buffer-accessor APIs across descriptors, atomic models, fitting and model layers; extends freeze to JIT-wrap buffer getters and accept do_atomic_virial; implements dual static (.pd) and predictor (.json) inference paths in DeepEval; refactors PD C++ nlist flattening; updates tests, typing, and docs. Changes
Sequence Diagram(s)sequenceDiagram autonumber participant U as User/CLI participant Freeze as freeze() participant Model as ModelModule participant J as paddle.jit.to_static U->>Freeze: freeze(model, output, head, do_atomic_virial) Freeze->>Model: build forward / forward_lower (InputSpec includes do_atomic_virial) Freeze->>Model: export static artifact (.pd / .json) alt buffer getters present loop each public get_buffer_* method Freeze->>J: to_static(method, input_spec=[], full_graph=True) J-->>Freeze: wrapped getter end end Freeze-->>U: saved frozen artifacts sequenceDiagram autonumber participant C as Caller participant DE as DeepEval participant PD as Static .pd Model participant PI as Predictor (.json) C->>DE: __init__(model_file, output_def, no_jit, ...) alt model_file is .pd and no_jit==False DE->>PD: load static .pd model PD-->>DE: static handles (buffers/outputs) else DE->>PI: create predictor for .json PI-->>DE: predictor handles (inputs/outputs) end C->>DE: eval(coords, box, types, fparam?, aparam?, spin?) alt static path DE->>PD: call forward / forward_lower PD-->>DE: outputs (tensors) else predictor path DE->>PI: set inputs & run predictor PI-->>DE: outputs end DE-->>C: numpy arrays (batched) Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews: pre_merge_checks: custom_checks: - name: "Undocumented Breaking Changes" mode: "warning" instructions: | Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🧹 Nitpick comments (12)
source/tests/pd/model/test_deeppot.py (4)
16-22: Consider refactoring imports by grouping utilities together.While the imports are functional, consider moving the PD-specific import
freezedirectly above the other PD importDeepPotfor better organization.from deepmd.infer.deep_pot import DeepPot as DeepPotUni -from deepmd.pd.entrypoints.main import ( - freeze, - get_trainer, -) from deepmd.pd.infer.deep_eval import ( DeepPot, ) +from deepmd.pd.entrypoints.main import ( + freeze, + get_trainer, +)
51-53: Consider using Path operations for file cleanup.Using
pathlib.Pathwould be more consistent with the rest of the code and safer for file operations.def tearDown(self) -> None: - for f in os.listdir("."): - if f in ["lcurve.out", self.input_json]: - os.remove(f) + for f in ["lcurve.out", self.input_json]: + if Path(f).exists(): + Path(f).unlink()
91-97: Add assertions for output values, not just shapes.While shape assertions are good, consider adding basic value checks to ensure the outputs are reasonable (e.g., energy/force values are not NaN or infinite).
ret = dp.eval(coord, cell, atype, atomic=True) e, f, v, ae, av = ret[0], ret[1], ret[2], ret[3], ret[4] self.assertEqual(e.shape, (1, 1)) self.assertEqual(f.shape, (1, 5, 3)) self.assertEqual(v.shape, (1, 9)) self.assertEqual(ae.shape, (1, 5, 1)) self.assertEqual(av.shape, (1, 5, 9)) + +# Verify outputs are valid numbers +self.assertFalse(np.isnan(e).any(), "Energy contains NaN values") +self.assertFalse(np.isnan(f).any(), "Force contains NaN values") +self.assertFalse(np.isinf(e).any(), "Energy contains infinite values") +self.assertFalse(np.isinf(f).any(), "Force contains infinite values")
110-117: Consider testing additional properties of type embeddings.While testing the shape and zero-padding is good, consider verifying that the non-zero elements of the type embedding have reasonable values.
def test_eval_typeebd(self) -> None: dp = DeepPot(str(self.model)) eval_typeebd = dp.eval_typeebd() self.assertEqual( eval_typeebd.shape, (len(self.config["model"]["type_map"]) + 1, 8) ) np.testing.assert_allclose(eval_typeebd[-1], np.zeros_like(eval_typeebd[-1])) + # Verify non-zero embeddings have reasonable values + for i in range(len(self.config["model"]["type_map"])): + self.assertFalse(np.allclose(eval_typeebd[i], 0), + f"Type embedding {i} should not be all zeros")deepmd/pd/infer/deep_eval.py (8)
112-114: Remove unusedno_jitandkwargsparameters.The
no_jitparameter in line 112 andkwargsin line 114 are not used in the__init__method.def __init__( self, model_file: str, output_def: ModelOutputDef, *args: Any, auto_batch_size: Union[bool, int, AutoBatchSize] = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, head: Optional[Union[str, int]] = None, - no_jit: bool = False, - **kwargs: Any, ) -> None:
191-195: Simplify the_has_spinattribute initialization.The conditional logic for
_has_spincan be more concise by directly checking the instance type.-self._has_spin = ( - getattr(self.dp.model["Default"], "has_spin", False) - if isinstance(self.dp, ModelWrapper) - else False -) +if isinstance(self.dp, ModelWrapper): + self._has_spin = getattr(self.dp.model["Default"], "has_spin", False) +else: + self._has_spin = False
260-261: Handle the missingget_var_namecase more gracefully.When
self.get_var_name()is called but raisesNotImplementedError, the error in model type determination could be clearer.elif "wfc" in model_output_type: return DeepWFC -elif self.get_var_name() in model_output_type: - return DeepProperty +else: + try: + var_name = self.get_var_name() + if var_name in model_output_type: + return DeepProperty + except NotImplementedError: + pass else: raise RuntimeError("Unknown model type")
488-490: Ensure dtype consistency foratype_handle.The dtype conversion for atom types uses a complex nested dictionary lookup. Consider simplifying for clarity.
-self.atype_handle.copy_from_cpu( - atom_types.astype( - NP_PRECISION_DICT[RESERVED_PRECISION_DICT[paddle.int64]] - ) -) +# Atom types should always be int32 for inference +self.atype_handle.copy_from_cpu(atom_types.astype(np.int32))
495-499: Consider implementing fparam and aparam support for static models.The static model path raises
NotImplementedErrorfor fparam and aparam inputs. This limits functionality compared to the dynamic path.Would you like me to help implement fparam and aparam support for static models, or should we track this as a known limitation in the documentation?
577-661: Refactor duplicate code between_eval_modeland_eval_model_spin.There's significant code duplication between these two methods. Consider extracting common logic.
The spin evaluation method shares much of its logic with the regular evaluation. Consider extracting a common helper:
def _prepare_inputs( self, coords: np.ndarray, cells: Optional[np.ndarray], atom_types: np.ndarray, fparam: Optional[np.ndarray], aparam: Optional[np.ndarray], spins: Optional[np.ndarray] = None, ) -> tuple: """Prepare inputs for model evaluation.""" nframes = coords.shape[0] if len(atom_types.shape) == 1: natoms = len(atom_types) atom_types = np.tile(atom_types, nframes).reshape([nframes, -1]) else: natoms = len(atom_types[0]) # Common coordinate preparation coord_input = paddle.to_tensor( coords.reshape([nframes, natoms, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE, ) type_input = paddle.to_tensor(atom_types, dtype=paddle.int64, place=DEVICE) # Optional inputs box_input = None if cells is None else paddle.to_tensor( cells.reshape([nframes, 3, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE, ) # Handle spin if provided spin_input = None if spins is None else paddle.to_tensor( spins.reshape([nframes, natoms, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE, ) # Parameter inputs fparam_input = None if fparam is None else to_paddle_tensor( fparam.reshape([nframes, self.get_dim_fparam()]) ) aparam_input = None if aparam is None else to_paddle_tensor( aparam.reshape([nframes, natoms, self.get_dim_aparam()]) ) return (nframes, natoms, coord_input, type_input, box_input, spin_input, fparam_input, aparam_input)
683-685: Document the shape calculation for Hessian outputs.The Hessian shape
[nframes, 3 * natoms, 3 * natoms]is correct but the commented line suggests uncertainty. Consider adding a clarifying comment.elif odef.category == OutputVariableCategory.DERV_R_DERV_R: + # Hessian matrix: second derivative of energy w.r.t. positions + # Shape is [nframes, 3*natoms, 3*natoms] for the full Hessian return [nframes, 3 * natoms, 3 * natoms] - # return [nframes, *odef.shape, 3 * natoms, 3 * natoms]
813-815: Clarify the model selection logic ineval_descriptor.The conditional logic for selecting between
ModelWrapperand raw model could be clearer with a comment.+# Select the appropriate model: unwrap ModelWrapper if present, otherwise use raw model model = ( self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
deepmd/pd/entrypoints/main.py(4 hunks)deepmd/pd/infer/deep_eval.py(20 hunks)deepmd/pd/model/atomic_model/base_atomic_model.py(2 hunks)deepmd/pd/model/atomic_model/dp_atomic_model.py(2 hunks)deepmd/pd/model/descriptor/dpa1.py(2 hunks)deepmd/pd/model/descriptor/dpa2.py(2 hunks)deepmd/pd/model/descriptor/dpa3.py(2 hunks)deepmd/pd/model/descriptor/repflows.py(3 hunks)deepmd/pd/model/descriptor/repformers.py(2 hunks)deepmd/pd/model/descriptor/se_a.py(4 hunks)deepmd/pd/model/descriptor/se_atten.py(2 hunks)deepmd/pd/model/descriptor/se_t_tebd.py(4 hunks)deepmd/pd/model/task/fitting.py(4 hunks)doc/third-party/ase.md(2 hunks)source/api_cc/src/DeepPotPD.cc(9 hunks)source/tests/pd/model/test_deeppot.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/model/descriptor/repflows.pydeepmd/pd/model/descriptor/dpa2.pydeepmd/pd/model/atomic_model/dp_atomic_model.pydeepmd/pd/model/descriptor/repformers.pydeepmd/pd/model/descriptor/dpa3.pydeepmd/pd/model/descriptor/se_atten.pydeepmd/pd/model/descriptor/dpa1.pydeepmd/pd/model/task/fitting.pydeepmd/pd/model/descriptor/se_t_tebd.pydeepmd/pd/model/atomic_model/base_atomic_model.pydeepmd/pd/model/descriptor/se_a.pydeepmd/pd/entrypoints/main.pysource/tests/pd/model/test_deeppot.pydeepmd/pd/infer/deep_eval.py
🧠 Learnings (3)
📓 Common learnings
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-10-08T15:32:11.479Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. 📚 Learning: 2024-09-19T04:25:12.408Z
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Applied to files:
source/tests/pd/model/test_deeppot.py
📚 Learning: 2025-09-18T11:37:10.520Z
Learnt from: CR PR: deepmodeling/deepmd-kit#0 File: AGENTS.md:0-0 Timestamp: 2025-09-18T11:37:10.520Z Learning: Applies to source/tests/tf/test_dp_test.py : Keep the core TensorFlow test `source/tests/tf/test_dp_test.py` passing; use it for quick validation Applied to files:
source/tests/pd/model/test_deeppot.py
🧬 Code graph analysis (12)
deepmd/pd/model/descriptor/dpa2.py (6)
deepmd/pd/model/descriptor/dpa3.py (2)
get_rcut(220-224)get_rcut_smth(226-230)deepmd/pd/model/descriptor/se_a.py (4)
get_rcut(120-122)get_rcut(515-519)get_rcut_smth(124-126)get_rcut_smth(521-525)deepmd/pd/model/descriptor/se_t_tebd.py (4)
get_rcut(190-192)get_rcut(617-621)get_rcut_smth(194-196)get_rcut_smth(623-627)deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_rcut(148-152)deepmd/pd/model/descriptor/repformers.py (2)
get_rcut(320-324)get_rcut_smth(326-330)deepmd/pd/model/descriptor/se_atten.py (2)
get_rcut(275-279)get_rcut_smth(281-285)
deepmd/pd/model/atomic_model/dp_atomic_model.py (7)
deepmd/pd/model/descriptor/dpa2.py (2)
get_rcut(330-334)get_sel(346-348)deepmd/pd/model/descriptor/dpa3.py (2)
get_rcut(220-224)get_sel(236-240)deepmd/pd/model/descriptor/se_a.py (4)
get_rcut(120-122)get_rcut(515-519)get_sel(132-134)get_sel(531-533)deepmd/pd/model/descriptor/repformers.py (2)
get_rcut(320-324)get_sel(336-338)deepmd/pd/model/descriptor/se_atten.py (2)
get_rcut(275-279)get_sel(291-293)deepmd/pd/model/descriptor/repflows.py (2)
get_rcut(356-360)get_sel(370-374)deepmd/dpmodel/atomic_model/dp_atomic_model.py (2)
get_rcut(62-64)get_sel(66-68)
deepmd/pd/model/descriptor/repformers.py (3)
deepmd/pd/model/descriptor/se_a.py (2)
get_rcut_smth(124-126)get_rcut_smth(521-525)deepmd/pd/model/descriptor/se_atten.py (1)
get_rcut_smth(281-285)deepmd/pd/model/descriptor/descriptor.py (1)
get_rcut_smth(63-65)
deepmd/pd/model/descriptor/dpa3.py (5)
deepmd/pd/model/descriptor/se_a.py (7)
get_rcut(120-122)get_rcut(515-519)get_rcut_smth(124-126)get_rcut_smth(521-525)get_sel(132-134)get_sel(531-533)get_type_map(140-144)deepmd/pd/model/descriptor/se_t_tebd.py (7)
get_rcut(190-192)get_rcut(617-621)get_rcut_smth(194-196)get_rcut_smth(623-627)get_sel(202-204)get_sel(633-635)get_type_map(210-214)deepmd/pd/model/descriptor/repflows.py (3)
get_rcut(356-360)get_rcut_smth(362-364)get_sel(370-374)deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_type_map(128-132)source/api_cc/src/DeepPotPD.cc (2)
get_type_map(659-668)get_type_map(659-659)
deepmd/pd/model/descriptor/se_atten.py (3)
deepmd/pd/model/descriptor/dpa1.py (1)
get_rcut_smth(326-328)deepmd/dpmodel/descriptor/dpa1.py (2)
get_rcut_smth(348-350)get_rcut_smth(799-801)deepmd/pd/model/descriptor/descriptor.py (1)
get_rcut_smth(63-65)
deepmd/pd/model/task/fitting.py (4)
deepmd/pd/infer/deep_eval.py (1)
get_dim_aparam(218-222)deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_dim_aparam(381-383)deepmd/dpmodel/fitting/general_fitting.py (1)
get_dim_aparam(233-235)deepmd/pt/model/task/fitting.py (1)
get_dim_aparam(450-452)
deepmd/pd/model/descriptor/se_t_tebd.py (3)
deepmd/pd/model/descriptor/se_a.py (2)
get_rcut_smth(124-126)get_rcut_smth(521-525)deepmd/pd/model/descriptor/repformers.py (1)
get_rcut_smth(326-330)deepmd/pd/model/descriptor/se_atten.py (1)
get_rcut_smth(281-285)
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
source/api_cc/src/common.cc (2)
split(36-47)split(36-37)
deepmd/pd/model/descriptor/se_a.py (5)
deepmd/pd/model/descriptor/dpa2.py (1)
get_rcut_smth(336-340)deepmd/pd/model/descriptor/dpa3.py (1)
get_rcut_smth(226-230)deepmd/pd/model/descriptor/se_t_tebd.py (2)
get_rcut_smth(194-196)get_rcut_smth(623-627)deepmd/pd/model/descriptor/repformers.py (1)
get_rcut_smth(326-330)deepmd/pd/model/descriptor/se_atten.py (1)
get_rcut_smth(281-285)
source/api_cc/src/DeepPotPD.cc (2)
source/api_cc/src/common.cc (1)
datype(753-753)source/api_cc/src/DeepPot.cc (16)
compute(80-93)compute(80-87)compute(96-107)compute(96-103)compute(109-116)compute(118-125)compute(127-134)compute(136-143)compute(146-162)compute(146-156)compute(165-179)compute(165-175)compute(181-191)compute(193-203)compute(205-215)compute(217-227)
source/tests/pd/model/test_deeppot.py (2)
deepmd/pd/entrypoints/main.py (3)
main(559-607)freeze(341-442)get_trainer(82-191)deepmd/pd/infer/deep_eval.py (7)
eval(309-383)get_type_map(208-210)get_ntypes(204-206)get_dim_fparam(212-216)get_dim_aparam(218-222)model_type(245-263)eval_typeebd(689-717)
deepmd/pd/infer/deep_eval.py (11)
deepmd/dpmodel/output_def.py (3)
ModelOutputDef(274-341)OutputVariableCategory(130-148)OutputVariableDef(151-239)deepmd/jax/infer/deep_eval.py (7)
DeepEval(60-434)get_model(426-434)eval(189-252)get_dim_fparam(143-145)get_dim_aparam(147-149)model_type(152-166)_get_output_shape(398-420)deepmd/dpmodel/infer/deep_eval.py (7)
DeepEval(57-413)get_model(405-413)eval(168-231)get_dim_fparam(118-120)get_dim_aparam(122-124)model_type(131-145)_get_output_shape(375-399)deepmd/pt/infer/deep_eval.py (11)
DeepEval(85-841)get_model(723-731)eval(306-380)get_dim_fparam(213-215)get_dim_aparam(217-219)get_intensive(229-230)get_var_name(232-239)model_type(242-260)_get_output_shape(623-647)_eval_model_spin(537-621)eval_fitting_last_layer(788-841)deepmd/tf/infer/deep_eval.py (8)
DeepEval(62-1137)get_model(1129-1137)eval(695-777)get_dim_fparam(646-648)get_dim_aparam(650-652)model_type(269-290)model_type(1205-1212)_get_output_shape(1004-1023)deepmd/infer/deep_property.py (1)
DeepProperty(21-164)deepmd/pd/model/network/network.py (1)
TypeEmbedNetConsistent(127-341)deepmd/utils/econf_embd.py (1)
sort_element_type(273-283)deepmd/utils/model_branch_dict.py (1)
get_model_dict(12-42)deepmd/pd/model/atomic_model/dp_atomic_model.py (3)
get_dim_fparam(377-379)get_dim_aparam(381-383)eval_fitting_last_layer(136-138)deepmd/pd/train/wrapper.py (1)
ModelWrapper(22-215)
🪛 GitHub Actions: Test Python
deepmd/pd/model/descriptor/dpa2.py
[error] 269-269: Failed to initialize buffer 'buffer_type_map' due to NoneType object, indicating missing or undefined 'type_map'.
deepmd/pd/model/atomic_model/dp_atomic_model.py
[error] 70-70: TypeError: unsupported operand type(s) for |: 'type' and 'types.GenericAlias'.
[error] 70-70: Syntax incompatible: using Python 3.10+ syntax with '|' in type hints, while environment does not support it.
[error] 70-70: TypeError at init: using unsupported Python syntax for type annotations ('|'), likely need to use 'Union' from 'typing'.
[error] 70-70: TypeError: unsupported operand type(s) for |: 'type' and 'types.GenericAlias'.
deepmd/pd/model/descriptor/dpa3.py
[error] 181-181: TypeError: 'NoneType' object is not iterable
deepmd/pd/model/descriptor/se_atten.py
[error] 156-156: Failed to initialize buffer 'buffer_type_map' because 'type_map' is None, indicating missing or undefined 'type_map'.
deepmd/pd/model/descriptor/dpa1.py
[error] 99-99: Failed to initialize buffer 'buffer_type_map' due to NoneType object, indicating missing or undefined 'type_map'.
deepmd/pd/model/task/fitting.py
[error] 434-434: Failed to deserialize data into 'Fitting' class, caused by 'type_map' being None.
deepmd/pd/model/descriptor/se_t_tebd.py
[error] 169-169: TypeError in init: 'NoneType' is not iterable when initializing buffer_type_map with type_one_side=None.
[error] 169-169: TypeError: ord() expected a character, string of length 2 found when attempting to convert 'type_map' which is None to tensor.
[error] 400-400: TypeError: 'NoneType' object is not iterable during deserialization.
[error] 400-400: TypeError during deserialization: 'NoneType' misses 'type_map' which is required.
[error] 400-400: TypeError: 'NoneType' during deserialization: 'type_map' was not specified.
deepmd/pd/model/descriptor/se_a.py
[error] 99-99: TypeError in init: 'NoneType' is not iterable when creating buffer_type_map with type_one_side=None.
[error] 183-183: TypeError: ord() expected a character, but string of length 2 found in buffer_type_map initialization when 'type_map' is None.
[error] 370-370: TypeError: 'NoneType' object is not iterable during deserialization.
[error] 370-370: TypeError during deserialization: 'NoneType' missing 'type_map' which is required.
[error] 370-370: TypeError: 'NoneType' during deserialization, missing 'type_map'.
🪛 Ruff (0.13.1)
deepmd/pd/infer/deep_eval.py
112-112: Unused method argument: no_jit
(ARG002)
113-113: Unused method argument: kwargs
(ARG002)
238-238: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
715-715: Avoid specifying long messages outside the exception class
(TRY003)
🔇 Additional comments (39)
source/api_cc/src/DeepPotPD.cc (8)
386-389: Type conversion from int to int64_t for atypeThe conversion to 64-bit integers looks correct for compatibility with Paddle's expected input types.
457-457: Replaced nlist construction withfillNlistTensorUsing the new helper function to populate the nlist tensor is a good refactoring that centralizes the tensor filling logic.
459-464: fparam and aparam tensor handling updatedThe change from using "box" input handles to explicit "fparam" and "aparam" handles better reflects the actual parameter purpose and improves code clarity.
Also applies to: 465-471
516-530: Template parameter names standardizedThe renaming from
dforce/dcoordtoforce/coordin template specializations improves consistency with the generic template declarations and matches the naming convention in DeepPot.cc.Also applies to: 531-545
579-583: Consistent fparam/aparam handling in non-ghost compute pathThe changes correctly use explicit input handles for fparam and aparam instead of the previous "box" workaround, maintaining consistency with the ghost atom compute path.
Also applies to: 585-590
632-642: Template instantiation parameter names alignedThe template instantiation parameter renaming maintains consistency with the updated method signatures. This improves code readability without affecting functionality.
Also applies to: 644-655
394-394: Message passing condition simplifiedThe condition change from
(do_message_passing == 1 && nghost > 0)to just(do_message_passing)removes the ghost check. This could lead to unnecessary operations when there are no ghost atoms.Is the removal of the
nghost > 0check intentional? If message passing operations are expensive, you might want to keep the optimization:- if (do_message_passing) { + if (do_message_passing && nghost > 0) {
391-391: Parameter change incopy_from_nlistcallThe change from using
nalltonall - nghostfor thecopy_from_nlistcall needs verification to ensure it correctly handles the local atom count.doc/third-party/ase.md (2)
5-5: LGTM! Excellent documentation enhancement.The addition of explicit backend support information helps users understand compatibility and choose appropriate backends for their needs.
10-92: Well-structured documentation with proper tab organization.The tab-based organization clearly separates usage examples for different backends. The code examples are consistent and demonstrate the correct file extensions (.pb for TensorFlow, .pth for PyTorch, .json for Paddle).
deepmd/pd/model/descriptor/dpa2.py (1)
332-340: Consistent dynamic/static mode handling pattern.The pattern of checking
paddle.in_dynamic_mode()to decide between instance variables and buffers is correctly implemented and consistent with the broader PR changes.deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
150-158: Consistent dynamic/static mode handling.The implementation correctly follows the established pattern for switching between dynamic and static mode data access.
deepmd/pd/model/descriptor/repflows.py (2)
248-250: LGTM! Proper buffer registration.The buffer registration is correctly implemented and follows the established pattern.
358-374: Consistent dynamic/static mode handling.The implementation correctly follows the pattern used across the codebase for switching between dynamic and static mode data access.
deepmd/pd/model/descriptor/se_atten.py (1)
156-158: LGTM! Proper buffer registration.Buffer registration is correctly implemented for the cut-off radius parameters.
deepmd/pd/model/descriptor/repformers.py (2)
213-215: LGTM! Proper buffer registration.Buffer registration follows the correct pattern established in other files.
322-330: Same potential return type issue as se_atten.py.The return type should be verified to ensure it returns
floatvalues, not tensors.Run the same verification script as above to ensure consistent behavior across all descriptor implementations.
deepmd/pd/entrypoints/main.py (3)
341-345: LGTM! Good addition of thedo_atomic_virialparameterThe addition of
do_atomic_virialas a configurable parameter provides better control over atomic virial computation during model freezing.
413-432: LGTM! Excellent addition of static method conversionsThe conversion of getter methods to static using
paddle.jit.to_staticis well-implemented and covers all essential model metadata accessors. This enables proper functionality in static graph mode.
373-378: Verify atype dtype change from int32 → int64Automated repository searches produced no matches; unable to confirm consistency. Manually verify that changing atype to int64 in deepmd/pd/entrypoints/main.py (lines 373–378 and 394–400) is safe for all consumers — Python callers, C++ bindings, model serialization, and any downstream code assuming int32.
deepmd/pd/model/task/fitting.py (4)
251-257: LGTM! Buffer registration for fparam and aparam dimensionsThe registration of
buffer_numb_fparamandbuffer_numb_aparamas persistent buffers correctly enables access to these parameters in static mode.
440-450: LGTM! Consistent implementation of static mode supportThe implementation of
get_dim_fparamandget_dim_aparamcorrectly returns buffer values in non-dynamic mode while preserving dynamic mode behavior.
266-269: Fix potential NoneType error whentype_mapis NoneSimilar to the issue in dpa3.py, the
type_mapcan beNonebut the code attempts to iterate over it when creatingbuffer_type_map.Apply this fix:
- self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in self.type_map]), - ) + if self.type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(self.type_map)]), + ) + else: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([], dtype=paddle.int32), + )
469-473: Handle None case inget_type_mapmethodThe method should handle the case when
type_mapisNoneto avoid errors during static mode execution.Apply this fix:
def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" if paddle.in_dynamic_mode(): - return self.type_map - return "".join([chr(x) for x in self.buffer_type_map.numpy()]).split(" ") + return self.type_map if self.type_map is not None else [] + if self.buffer_type_map.numel() == 0: + return [] + return "".join([chr(x) for x in self.buffer_type_map.numpy()]).split(" ")deepmd/pd/model/descriptor/se_a.py (3)
446-448: LGTM! Proper buffer registration for static mode supportThe registration of
buffer_rcutandbuffer_rcut_smthas persistent buffers and their usage in the corresponding getter methods correctly implements static mode support for the descriptor block.Also applies to: 515-525
98-100: Fix NoneType error whentype_mapis NoneThe
type_mapparameter can beNonebut the code attempts to iterate over it when creatingbuffer_type_map.Apply this fix:
- self.register_buffer( - "buffer_type_map", paddle.to_tensor([ord(c) for c in type_map]) - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", paddle.to_tensor([ord(c) for c in " ".join(type_map)]) + ) + else: + self.register_buffer( + "buffer_type_map", paddle.to_tensor([], dtype=paddle.int32) + )
140-144: Handle None case inget_type_mapmethodThe method should handle the case when
type_mapisNoneto avoid errors during static mode execution.Apply this fix:
def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" if paddle.in_dynamic_mode(): - return self.type_map - return "".join([chr(x) for x in self.buffer_type_map.numpy()]).split(" ") + return self.type_map if self.type_map is not None else [] + if self.buffer_type_map.numel() == 0: + return [] + return "".join([chr(x) for x in self.buffer_type_map.numpy()]).split(" ")deepmd/pd/model/descriptor/se_t_tebd.py (3)
543-545: LGTM! Consistent implementation of static mode buffersThe registration of
buffer_rcutandbuffer_rcut_smthand their usage in the corresponding getter methods correctly implements static mode support, consistent with other descriptor implementations.Also applies to: 617-627
168-170: Fix NoneType error whentype_mapis NoneThe
type_mapparameter can beNonebut the code attempts to iterate over it when creatingbuffer_type_map.Apply this fix:
- self.register_buffer( - "buffer_type_map", paddle.to_tensor([ord(c) for c in type_map]) - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", paddle.to_tensor([ord(c) for c in " ".join(type_map)]) + ) + else: + self.register_buffer( + "buffer_type_map", paddle.to_tensor([], dtype=paddle.int32) + )
210-214: Handle None case inget_type_mapmethodThe method should handle the case when
type_mapisNoneto avoid errors during static mode execution.Apply this fix:
def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" if paddle.in_dynamic_mode(): - return self.type_map - return "".join([chr(x) for x in self.buffer_type_map.numpy()]).split(" ") + return self.type_map if self.type_map is not None else [] + if self.buffer_type_map.numel() == 0: + return [] + return "".join([chr(x) for x in self.buffer_type_map.numpy()]).split(" ")source/tests/pd/model/test_deeppot.py (4)
45-47: Use context manager when switching device.The
paddle.device("cpu")context should be properly managed to ensure device state is restored after use.The code uses
with paddle.device("cpu"):correctly, so this is actually fine.
56-90: LGTM!The test setup with realistic water molecule coordinates and cell parameters is well-structured for testing the DeepPot inference functionality.
119-130: Skip reason is clear and appropriate.The skip decorator with a clear reason for why frozen models don't support attribute access is helpful for test maintainability.
134-138: Mock patches may not fully disable CUDA.The comment correctly notes that these mocks don't actually disable CUDA. Consider adding a more explicit check or warning about the limitations of this test approach.
This test is trying to ensure device compatibility, but the approach may not be fully effective. The comment acknowledges this limitation, which is good for transparency.
deepmd/pd/infer/deep_eval.py (5)
151-152: Improve path manipulation for.jsonsuffix removal.Using string splitting with
split(".json")[0]is fragile and could break if the filename contains.jsonelsewhere. Consider usingpathlib.Pathoros.path.splitextfor more robust path manipulation.-self.dp = paddle.jit.load(self.model_path.split(".json")[0]) +from pathlib import Path +self.dp = paddle.jit.load(str(Path(self.model_path).with_suffix("")))
414-415: LGTM! Hessian support correctly added.The addition of
OutputVariableCategory.DERV_R_DERV_Rto the request definitions enables proper Hessian computation support.
710-717: LGTM! Type embedding evaluation correctly implemented.The implementation properly collects and concatenates type embeddings from all
TypeEmbedNetConsistentsublayers, with appropriate error handling when no embeddings are found.
161-161: Document or make configurable the hardcoded GPU configuration values.The hardcoded values
4096, 0for GPU configuration should be made configurable or at least documented. Magic numbers like these reduce maintainability and may not be appropriate for all deployment scenarios.Consider making these values configurable:
+# GPU memory pool size in MB and device ID +GPU_MEMORY_POOL_SIZE = int(os.environ.get("DEEPMD_GPU_MEMORY_POOL_SIZE", "4096")) +GPU_DEVICE_ID = int(os.environ.get("DEEPMD_GPU_DEVICE_ID", "0")) + config.enable_custom_passes( ["add_shadow_output_after_dead_parameter_pass"], True ) -config.enable_use_gpu(4096, 0) +config.enable_use_gpu(GPU_MEMORY_POOL_SIZE, GPU_DEVICE_ID)
214-222: Handle the static model case properly for parameter dimensions.The methods
get_dim_fparam()andget_dim_aparam()return hardcoded values in the static model path, but this may not be correct for all model types. These should properly delegate to the model's actual implementation.The implementation is actually correct - it delegates to
self.dp.get_dim_fparam()andself.dp.get_dim_aparam()for static models. The previous review comment appears to be outdated.
HydrogenSulfate force-pushed the pd_deepeval branch from fdb796d to 80a4c5f Compare ![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (9)
deepmd/pd/model/descriptor/dpa3.py (1)
317-345: Update buffers when type_map changes.
change_type_map()updates fields but leavesbuffer_ntypes/buffer_type_mapstale, breaking static-mode consumers.self.type_embedding.change_type_map(type_map=type_map) self.exclude_types = map_pair_exclude_types(self.exclude_types, remap_index) - self.ntypes = len(type_map) + self.ntypes = len(type_map) repflow = self.repflows @@ repflow["dstd"] = repflow["dstd"][remap_index] + + # keep static buffers in sync + try: + self.buffer_ntypes = paddle.to_tensor(self.ntypes, dtype="int64") + except Exception: + pass + new_btm = paddle.to_tensor( + [ord(c) for c in " ".join(self.type_map)] if self.type_map is not None else [], + dtype="int32", + ) + try: + self.buffer_type_map = new_btm + except Exception: + # fallback when direct assignment is constrained + paddle.assign(new_btm, self.buffer_type_map)deepmd/pd/model/atomic_model/base_atomic_model.py (3)
133-136: Support static-mode get_type_map and empty buffers.Expose
list[str]consistently and avoid.numpy()on empty buffers.- return self.type_map + if paddle.in_dynamic_mode(): + return self.type_map if self.type_map is not None else [] + if self.buffer_type_map.numel() == 0: + return [] + return "".join([chr(int(x)) for x in self.buffer_type_map.numpy()]).split(" ")
309-339: Keep static buffers in sync in change_type_map().After remapping, update
buffer_ntypesandbuffer_type_mapfor static inference.self.type_map = type_map @@ - if has_new_type: + if has_new_type: extend_shape = [ self.out_bias.shape[0], len(type_map), *list(self.out_bias.shape[2:]), ] @@ self.out_bias = self.out_bias[:, remap_index, :] self.out_std = self.out_std[:, remap_index, :] + + # sync buffers for static mode + self.ntypes = len(type_map) + self.buffer_ntypes = paddle.to_tensor(self.ntypes, dtype="int64") + self.buffer_type_map = paddle.to_tensor( + [ord(c) for c in " ".join(self.type_map)], dtype="int32" + )
1-10: Run ruff and fix lint/format errors before merge.ruff reported multiple violations (e.g. B006 mutable-default-args; ARG*/RUF*/F841 unused args/vars). Fix the failures (notably deepmd/pd/model/atomic_model/base_atomic_model.py:78,79,159,310,440), then run: ruff check . && ruff format . and ensure CI passes.
source/api_cc/src/DeepPotPD.cc (2)
394-444: Guard message-passing tensors when no swaps; fix sendlist shapeAvoid constructing zero-length tensors unless guaranteed supported; and sendlist currently reshaped to total_send but populated with nswap pointers — mismatch.
- if (do_message_passing) { + if (do_message_passing && lmp_list.nswap > 0) { @@ - int total_send = - std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0); - sendlist_tensor->Reshape({total_send}); + // pass one pointer per swap + sendlist_tensor->Reshape({nswap});
584-589: Validate aparam divisibility before reshape (dynamic path)Protect against non-divisible aparam.size() by natoms to prevent invalid shapes.
- aparam_tensor->Reshape( - {1, natoms, static_cast<int>(aparam.size()) / natoms}); + if (aparam.size() % static_cast<size_t>(natoms) != 0) { + throw deepmd::deepmd_exception("aparam size is not divisible by natoms"); + } + aparam_tensor->Reshape( + {1, natoms, static_cast<int>(aparam.size() / natoms)});deepmd/pd/entrypoints/main.py (1)
368-381: InputSpec replaced with a boolean: to_static will receive an invalid input_specPassing
do_atomic_virial(bool) insideinput_specis wrong; it must be anInputSpecor omitted to rely on the function default. This will fail at export/runtime.Apply:
- # InputSpec([], dtype="bool", name="do_atomic_virial"), # do_atomic_virial - do_atomic_virial, # do_atomic_virial + # Keep defaulted inside the function; don't expose as input here (or add an InputSpec and plumb it in predictor). + # InputSpec([], dtype="bool", name="do_atomic_virial"), # do_atomic_virialIf you want it controllable at inference time, keep the
InputSpec([], dtype="bool")and add a matching predictor input handle indeep_eval.py.deepmd/pd/infer/deep_eval.py (2)
819-835: Descriptor eval not supported for static modelsStatic
.jsonobjects won’t haveset_eval_descriptor_hook/eval_descriptor. Guard and raise a clear error for static mode.Apply:
- model = ( - self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp - ) + if not isinstance(self.dp, ModelWrapper): + raise NotImplementedError("eval_descriptor is not available for static (.json) models.") + model = self.dp.model["Default"]
836-890: fitting-before-last-layer: dynamic-only; guard static modeSame rationale as descriptor eval.
Apply:
- model = self.dp.model["Default"] + if not isinstance(self.dp, ModelWrapper): + raise NotImplementedError("eval_fitting_last_layer is not available for static (.json) models.") + model = self.dp.model["Default"]
🧹 Nitpick comments (16)
deepmd/pd/model/descriptor/se_atten.py (1)
301-301: Fix return type annotation mismatch in non-dynamic mode.The method is annotated to return
int, but in non-dynamic mode it returns a tensor. Ensure consistent return type by calling.item().Apply this diff to fix the type annotation mismatch:
def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes.item()deepmd/pd/model/task/fitting.py (2)
443-445: Fix return type annotation mismatch in non-dynamic mode.The method is annotated to return
int, but in non-dynamic mode it returns a tensor. Ensure consistent return type by calling.item().Apply this diff to fix the type annotation mismatch:
def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" if paddle.in_dynamic_mode(): return self.numb_fparam - return self.buffer_numb_fparam + return self.buffer_numb_fparam.item()
449-451: Fix return type annotation mismatch in non-dynamic mode.The method is annotated to return
int, but in non-dynamic mode it returns a tensor. Ensure consistent return type by calling.item().Apply this diff to fix the type annotation mismatch:
def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" if paddle.in_dynamic_mode(): return self.numb_aparam - return self.buffer_numb_aparam + return self.buffer_numb_aparam.item()deepmd/pd/model/descriptor/se_t_tebd.py (2)
625-633: Fix return type annotation mismatches in non-dynamic mode.Both methods are annotated to return
float, but in non-dynamic mode they return tensors. Ensure consistent return types by calling.item().Apply this diff to fix the type annotation mismatches:
def get_rcut(self) -> float: """Returns the cut-off radius.""" if paddle.in_dynamic_mode(): return self.rcut - return self.buffer_rcut + return self.buffer_rcut.item() def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" if paddle.in_dynamic_mode(): return self.rcut_smth - return self.buffer_rcut_smth + return self.buffer_rcut_smth.item()
645-645: Fix return type annotation mismatch in non-dynamic mode.The method is annotated to return
int, but in non-dynamic mode it returns a tensor. Ensure consistent return type by calling.item().Apply this diff to fix the type annotation mismatch:
def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes.item()deepmd/pd/model/descriptor/dpa3.py (2)
114-122: Avoid mutable default args.
exclude_types: list[tuple[int, int]] = []is risky; useNoneand normalize.- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None, @@ - exclude_types=exclude_types, + exclude_types=[] if exclude_types is None else exclude_types,
239-244: Return type may vary in static mode.
get_sel()advertiseslist[int], but static path proxiesrepflows.get_sel()which can be a tensor; align typing or convert.deepmd/pd/model/atomic_model/base_atomic_model.py (1)
75-83: Avoid mutable default args.
atom_exclude_types=[]andpair_exclude_types=[]should useNonethen normalize.- atom_exclude_types: list[int] = [], - pair_exclude_types: list[tuple[int, int]] = [], + atom_exclude_types: Optional[list[int]] = None, + pair_exclude_types: Optional[list[tuple[int, int]]] = None, @@ - self.reinit_atom_exclude(atom_exclude_types) - self.reinit_pair_exclude(pair_exclude_types) + self.reinit_atom_exclude([] if atom_exclude_types is None else atom_exclude_types) + self.reinit_pair_exclude([] if pair_exclude_types is None else pair_exclude_types)source/api_cc/src/DeepPotPD.cc (1)
457-457: Remove unused variable do_atom_virial_tensorThis local is never used; drop it to avoid warnings.
- bool do_atom_virial_tensor = atomic;Also applies to: 591-591
deepmd/pd/entrypoints/main.py (1)
345-346: New freeze() arg added but not wired through CLI
do_atomic_virialis added tofreeze()but never propagated from CLI flags. Either thread a flag throughparse_argsor drop the parameter.deepmd/pd/infer/deep_eval.py (6)
112-115: Unused parameters: no_jit, kwargsThey are never used.
Apply:
- head: Optional[Union[str, int]] = None, - no_jit: bool = False, - **kwargs: Any, + head: Optional[Union[str, int]] = None,Or document/use them (e.g., gate JIT usage with
no_jit).
193-201: Avoid setattr/getattr with constant attribute names; simplify has_spin detectionCleaner and ruff‑compliant.
Apply:
- self._has_spin = ( - getattr(self.dp.model["Default"], "has_spin", False) - if isinstance(self.dp, ModelWrapper) - else False - ) - if callable(self._has_spin): - setattr(self, "_has_spin", self._has_spin()) + self._has_spin = False + if isinstance(self.dp, ModelWrapper): + has_spin_attr = getattr(self.dp.model["Default"], "has_spin", False) + self._has_spin = has_spin_attr() if callable(has_spin_attr) else bool(has_spin_attr)
216-224: Ensure int return types for static helpersStatic
self.dp.get_dim_fparam()/get_dim_aparam()may return Tensors; cast to int.Apply:
- if self.static_model: - return self.dp.get_dim_fparam() + if self.static_model: + val = self.dp.get_dim_fparam() + return int(val.numpy().item()) if hasattr(val, "numpy") else int(val)Repeat similarly for
get_dim_aparam().
237-245: Use direct attribute instead of getattr for get_var_nameSimplify and avoid ruff B009.
Apply:
- if hasattr(self.dp.model["Default"], "get_var_name") and callable( - getattr(self.dp.model["Default"], "get_var_name") - ): - return self.dp.model["Default"].get_var_name() + get_name = getattr(self.dp.model["Default"], "get_var_name", None) + if callable(get_name): + return get_name()
716-724: Type embedding extraction LGTM; minor: exception message styleImplementation is fine. Optional: shorten message to satisfy TRY003 if you enforce it.
738-747: get_model_size: static models unsupported
named_parameters()is unavailable on static graphs. Guard and raise.Apply:
- params_dict = dict(self.dp.named_parameters()) + if not isinstance(self.dp, ModelWrapper): + raise NotImplementedError("get_model_size is not available for static (.json) models.") + params_dict = dict(self.dp.named_parameters())
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (15)
deepmd/pd/entrypoints/main.py(4 hunks)deepmd/pd/infer/deep_eval.py(20 hunks)deepmd/pd/model/atomic_model/base_atomic_model.py(2 hunks)deepmd/pd/model/descriptor/dpa1.py(2 hunks)deepmd/pd/model/descriptor/dpa2.py(4 hunks)deepmd/pd/model/descriptor/dpa3.py(3 hunks)deepmd/pd/model/descriptor/repflows.py(3 hunks)deepmd/pd/model/descriptor/repformers.py(3 hunks)deepmd/pd/model/descriptor/se_a.py(6 hunks)deepmd/pd/model/descriptor/se_atten.py(4 hunks)deepmd/pd/model/descriptor/se_t_tebd.py(6 hunks)deepmd/pd/model/model/model.py(1 hunks)deepmd/pd/model/task/fitting.py(4 hunks)doc/third-party/ase.md(2 hunks)source/api_cc/src/DeepPotPD.cc(10 hunks)
🚧 Files skipped from review as they are similar to previous changes (6)
- deepmd/pd/model/descriptor/dpa1.py
- deepmd/pd/model/descriptor/se_a.py
- doc/third-party/ase.md
- deepmd/pd/model/descriptor/dpa2.py
- deepmd/pd/model/descriptor/repflows.py
- deepmd/pd/model/descriptor/repformers.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/model/task/fitting.pydeepmd/pd/model/descriptor/se_t_tebd.pydeepmd/pd/model/descriptor/se_atten.pydeepmd/pd/model/atomic_model/base_atomic_model.pydeepmd/pd/entrypoints/main.pydeepmd/pd/model/descriptor/dpa3.pydeepmd/pd/model/model/model.pydeepmd/pd/infer/deep_eval.py
🧠 Learnings (1)
📓 Common learnings
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-10-08T15:32:11.479Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. 🧬 Code graph analysis (8)
deepmd/pd/model/task/fitting.py (5)
deepmd/pd/infer/deep_eval.py (1)
get_dim_aparam(220-224)deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_dim_aparam(381-383)deepmd/dpmodel/fitting/general_fitting.py (1)
get_dim_aparam(233-235)deepmd/pt/model/task/fitting.py (1)
get_dim_aparam(450-452)deepmd/pd/model/model/make_model.py (1)
get_dim_aparam(528-530)
deepmd/pd/model/descriptor/se_t_tebd.py (3)
deepmd/pd/model/descriptor/se_a.py (2)
get_rcut_smth(126-128)get_rcut_smth(526-530)deepmd/pd/model/descriptor/repformers.py (1)
get_rcut_smth(330-334)deepmd/pd/model/descriptor/se_atten.py (1)
get_rcut_smth(285-289)
deepmd/pd/model/descriptor/se_atten.py (4)
deepmd/pd/model/descriptor/dpa1.py (1)
get_rcut_smth(328-330)deepmd/pd/model/descriptor/se_a.py (2)
get_rcut_smth(126-128)get_rcut_smth(526-530)deepmd/pd/model/descriptor/se_t_tebd.py (2)
get_rcut_smth(196-198)get_rcut_smth(629-633)deepmd/pd/model/descriptor/repformers.py (1)
get_rcut_smth(330-334)
deepmd/pd/model/atomic_model/base_atomic_model.py (4)
deepmd/pd/model/descriptor/dpa2.py (1)
get_ntypes(355-357)deepmd/pd/infer/deep_eval.py (1)
get_ntypes(206-208)deepmd/pd/model/descriptor/se_atten.py (1)
get_ntypes(299-301)deepmd/pd/model/model/model.py (1)
get_ntypes(53-58)
deepmd/pd/model/descriptor/dpa3.py (5)
deepmd/pd/model/descriptor/se_t_tebd.py (11)
get_rcut(192-194)get_rcut(623-627)get_rcut_smth(196-198)get_rcut_smth(629-633)get_nsel(200-202)get_nsel(635-637)get_sel(204-206)get_sel(639-641)get_ntypes(208-210)get_ntypes(643-645)get_type_map(212-216)deepmd/pd/model/descriptor/repflows.py (5)
get_rcut(359-363)get_rcut_smth(365-367)get_nsel(369-371)get_sel(373-377)get_ntypes(379-381)deepmd/pd/model/descriptor/se_atten.py (5)
get_rcut(279-283)get_rcut_smth(285-289)get_nsel(291-293)get_sel(295-297)get_ntypes(299-301)deepmd/pd/model/atomic_model/base_atomic_model.py (2)
get_ntypes(583-586)get_type_map(133-135)source/api_cc/src/DeepPotPD.cc (2)
get_type_map(658-667)get_type_map(658-658)
source/api_cc/src/DeepPotPD.cc (2)
source/api_cc/tests/test_deeppot_pd.cc (2)
ener(136-142)ener(136-140)source/api_cc/src/DeepPot.cc (16)
compute(80-93)compute(80-87)compute(96-107)compute(96-103)compute(109-116)compute(118-125)compute(127-134)compute(136-143)compute(146-162)compute(146-156)compute(165-179)compute(165-175)compute(181-191)compute(193-203)compute(205-215)compute(217-227)
deepmd/pd/model/model/model.py (11)
deepmd/pd/model/atomic_model/base_atomic_model.py (2)
get_type_map(133-135)get_ntypes(583-586)deepmd/pd/model/descriptor/dpa1.py (2)
get_type_map(344-348)get_ntypes(340-342)deepmd/pd/model/descriptor/dpa2.py (2)
get_type_map(359-361)get_ntypes(355-357)deepmd/pd/model/descriptor/dpa3.py (2)
get_type_map(249-253)get_ntypes(245-247)deepmd/pd/model/descriptor/se_a.py (3)
get_type_map(142-146)get_ntypes(138-140)get_ntypes(540-542)deepmd/pd/model/descriptor/se_t_tebd.py (3)
get_type_map(212-216)get_ntypes(208-210)get_ntypes(643-645)deepmd/pd/model/task/fitting.py (1)
get_type_map(470-474)deepmd/pd/infer/deep_eval.py (2)
get_type_map(210-212)get_ntypes(206-208)deepmd/pd/model/descriptor/repflows.py (1)
get_ntypes(379-381)deepmd/pd/model/descriptor/repformers.py (1)
get_ntypes(344-346)deepmd/pd/model/descriptor/se_atten.py (1)
get_ntypes(299-301)
deepmd/pd/infer/deep_eval.py (5)
deepmd/infer/deep_eval.py (20)
output_def(420-421)DeepEval(365-741)DeepEvalBackend(31-362)get_model(351-362)get_model(729-741)get_rcut(150-151)get_rcut(423-425)get_type_map(158-159)get_type_map(431-433)eval(102-147)get_dim_fparam(162-163)get_dim_fparam(435-437)get_dim_aparam(170-171)get_dim_aparam(443-445)has_default_fparam(165-167)has_default_fparam(439-441)model_type(302-303)get_has_spin(322-324)eval_fitting_last_layer(224-264)eval_fitting_last_layer(538-603)deepmd/pd/model/network/network.py (1)
TypeEmbedNetConsistent(127-341)deepmd/utils/econf_embd.py (1)
sort_element_type(273-283)deepmd/utils/model_branch_dict.py (1)
get_model_dict(12-42)deepmd/pd/train/wrapper.py (1)
ModelWrapper(22-215)
🪛 Ruff (0.13.1)
deepmd/pd/infer/deep_eval.py
112-112: Unused method argument: no_jit
(ARG002)
113-113: Unused method argument: kwargs
(ARG002)
199-199: Do not call setattr with a constant attribute value. It is not any safer than normal property access.
Replace setattr with assignment
(B010)
240-240: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
721-721: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
🔇 Additional comments (27)
deepmd/pd/model/descriptor/se_atten.py (3)
156-158: LGTM!The buffer registrations for
buffer_rcutandbuffer_rcut_smthare correctly implemented for static-mode compatibility.
190-192: LGTM!The buffer registration for
buffer_ntypesis correctly implemented using the appropriate int64 dtype for static-mode compatibility.
281-289: Use.item()on buffer tensors to return Python floats in static mode.The static-mode return paths still return paddle Tensors but the methods are annotated to return
float. Change the static-mode returns to return Python floats by calling.item()on the buffer tensors.Apply this diff to match the type annotation:
def get_rcut(self) -> float: """Returns the cut-off radius.""" if paddle.in_dynamic_mode(): return self.rcut - return self.buffer_rcut + return self.buffer_rcut.item() def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" if paddle.in_dynamic_mode(): return self.rcut_smth - return self.buffer_rcut_smth + return self.buffer_rcut_smth.item()deepmd/pd/model/task/fitting.py (3)
251-253: LGTM!The buffer registration for frame parameter dimension is correctly implemented for static-mode compatibility.
255-257: LGTM!The buffer registration for atomic parameter dimension is correctly implemented for static-mode compatibility.
266-270: LGTM!The conditional buffer registration for type_map is correctly implemented. The pattern of converting string characters to ord values and joining with spaces is consistent with similar implementations in other files.
deepmd/pd/model/descriptor/se_t_tebd.py (3)
168-172: LGTM!The conditional buffer registration for type_map follows the consistent pattern used in other files, converting string characters to ord values and joining with spaces.
545-547: LGTM!The buffer registrations for
buffer_rcutandbuffer_rcut_smthare correctly implemented for static-mode compatibility.
565-567: LGTM!The buffer registration for
buffer_ntypesis correctly implemented using the appropriate int64 dtype for static-mode compatibility.deepmd/pd/model/descriptor/dpa3.py (3)
214-216: LGTM: buffer_ntypes persisted for static mode.Int64 buffer for
ntypesis appropriate and consistent with other descriptors.Please confirm the C++ predictor path accepts int64 for
buffer_ntypesacross targets.
179-182: Handle None type_map and standardize dtype for buffer_type_map.Current code crashes when
type_mapis None and leaves dtype implicit. Fix and align with int32 used elsewhere/CC side.Apply:
- self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(self.type_map)]), - ) + if self.type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(self.type_map)], dtype="int32"), + ) + else: + self.register_buffer( + "buffer_type_map", + paddle.empty([0], dtype="int32"), + )
249-253: get_type_map returns a Tensor in static mode; return list[str] and handle empty.Docstring/typing promise
list[str], but static path returns a tensor. Decode buffer and guard empty.- if paddle.in_dynamic_mode(): - return self.type_map - return self.buffer_type_map + if paddle.in_dynamic_mode(): + return self.type_map if self.type_map is not None else [] + if self.buffer_type_map.numel() == 0: + return [] + return "".join([chr(int(x)) for x in self.buffer_type_map.numpy()]).split(" ")deepmd/pd/model/atomic_model/base_atomic_model.py (2)
583-587: LGTM: get_ntypes matches project pattern.Dynamic returns Python int; static returns buffer for graph tracing.
Please ensure downstream shape construction never passes a Tensor where a Python int is required in eager paths.
87-95: Fix None handling and dtype for buffer_type_map; prevent crash computing ntypes.
len(self.type_map)raises whentype_mapis None; also set explicit int32 for buffer.- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) - self.ntypes = len(self.type_map) - self.register_buffer( - "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)], dtype="int32"), + ) + else: + self.register_buffer( + "buffer_type_map", + paddle.empty([0], dtype="int32"), + ) + self.ntypes = len(self.type_map) if self.type_map is not None else 0 + self.register_buffer("buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64"))source/api_cc/src/DeepPotPD.cc (5)
456-457: Confirm nlist row uniformity post-shuffle/paddingfillNlistTensor assumes fixed nnei across rows. Ensure nlist_data.shuffle_exclude_empty + padding produce uniform row lengths (or update to pad here).
568-570: atype dtype widened to int64 — ensure model input expectationsforward path now feeds int64, while forward_lower path still feeds int32. Verify both model inputs’ dtypes; otherwise introduce a consistent dtype (or dual-path handling).
515-529: Template instantiation updates look consistentENERGYVTYPE vector instantiation aligns with upstream signature changes.
531-544: Template instantiation updates (float) look consistentNo issues spotted.
167-183: Fix ragged nlist handling and prevent shape/data mismatches in fillNlistTensornnei is derived via total_size / nloc without validating uniform row length; this can silently mis-shape the tensor or crash CopyFromCpu. Also prefer explicit int32 for Paddle int tensors and guard int casts. This was raised previously and remains unresolved.
Apply:
-void fillNlistTensor(const std::vector<std::vector<int>>& data, - std::unique_ptr<paddle_infer::Tensor>& flat_tensor) { - size_t total_size = 0; - for (const auto& row : data) { - total_size += row.size(); - } - std::vector<int> flat_data; - flat_data.reserve(total_size); - for (const auto& row : data) { - flat_data.insert(flat_data.end(), row.begin(), row.end()); - } - - int nloc = data.size(); - int nnei = nloc > 0 ? total_size / nloc : 0; - flat_tensor->Reshape({1, nloc, nnei}); - flat_tensor->CopyFromCpu(flat_data.data()); -} +void fillNlistTensor(const std::vector<std::vector<int>>& data, + std::unique_ptr<paddle_infer::Tensor>& flat_tensor) { + const int nloc = static_cast<int>(data.size()); + size_t total_size = 0; + const size_t expected_row = nloc > 0 ? data[0].size() : 0; + for (const auto& row : data) { + if (row.size() != expected_row) { + throw deepmd::deepmd_exception("Inconsistent row sizes in nlist data"); + } + total_size += row.size(); + } + std::vector<std::int32_t> flat_data; + flat_data.reserve(total_size); + for (const auto& row : data) { + flat_data.insert(flat_data.end(), row.begin(), row.end()); + } + int nnei = 0; + if (nloc > 0) { + if (expected_row > static_cast<size_t>(std::numeric_limits<int>::max())) { + throw deepmd::deepmd_exception("nnei exceeds int range"); + } + nnei = static_cast<int>(expected_row); + if (static_cast<size_t>(nloc) * static_cast<size_t>(nnei) != total_size) { + throw deepmd::deepmd_exception("Invalid nlist dimensions"); + } + } + flat_tensor->Reshape({1, nloc, nnei}); + flat_tensor->CopyFromCpu(flat_data.data()); +}deepmd/pd/infer/deep_eval.py (7)
152-157: Fragile path handling for .json
split(".json")[0]breaks for names containing.jsonearlier in the path.Apply:
- self.dp = paddle.jit.load(self.model_path.split(".json")[0]) + from pathlib import Path + base = str(Path(self.model_path).with_suffix("")) + self.dp = paddle.jit.load(base)
416-418: DERV_R_DERV_R inclusion looks goodExtending requested outputs to include Hessian category is correct.
556-566: Dynamic call returns tuple in some cases: OK, but consider guarding keysYou already unwrap tuples; fine. Ensure downstream expects dict keys consistent with
_OUTDEF_DP2BACKEND.
669-692: Hessian output shape handling LGTMShape
[nframes, 3*natoms, 3*natoms]is consistent with combined second derivatives.
753-768: Observed types API LGTMSorting via
sort_element_typeis appropriate.
2-2: Run ruff before commitPlease run:
- ruff check .
- ruff format .
to address small nits (ARG002, B009/B010, TRY003) flagged in the hints.
160-166: Hardcoded GPU config and custom pass
enable_use_gpu(4096, 0)and custom pass list are magic numbers/toggles that may not fit all envs; they can crash on CPU‑only hosts.
- Detect GPU availability and fall back to
disable_gpu(); make mem/idx configurable.- Guard custom pass behind availability/version checks.
Provide knobs via env vars or args.deepmd/pd/entrypoints/main.py (1)
389-412: Confirm boolean input_spec and atype dtype
- do_atomic_virial is provided as a boolean literal/variable in the to_static input_spec — decide whether this should be a runtime boolean input (replace with InputSpec([], dtype="bool", name="do_atomic_virial")) or an intentional constant (keep literal). (docs.deepmodeling.com)
- atype dtype differs: forward uses dtype="int64" while forward_lower uses dtype="int32" — confirm forward_lower truly expects int32; if not, change InputSpec([-1, -1], dtype="int64", name="atype") to align with forward. (docs.deepmodeling.com)
File: deepmd/pd/entrypoints/main.py (freeze → to_static for forward / forward_lower).
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pd/infer/deep_eval.py (1)
127-134: Fix multitask head index handlingUsing an int head currently ignores the provided index. Map the index to the correct key (with bounds check).
Apply:
- if isinstance(head, int): - head = model_keys[0] + if isinstance(head, int): + if head < 0 or head >= len(model_keys): + raise IndexError(f"Head index {head} out of range [0, {len(model_keys)-1}]") + head = model_keys[head]
🧹 Nitpick comments (6)
deepmd/pd/infer/deep_eval.py (6)
2-14: Use pathlib and add os/env support for inference configPreload pathlib/os for robust paths and config via env/kwargs.
Apply:
import logging +import os +from pathlib import Path @@ -from paddle import inference as paddle_inference +from paddle import inference as paddle_inference
226-236: has_default_fparam: support static modelsUse backend method when frozen.
Apply:
def has_default_fparam(self) -> bool: """Check if the model has default frame parameters.""" - try: - return self.dp.model["Default"].has_default_fparam() - except AttributeError: - # for compatibility with old models - return False + if self.static_model and hasattr(self.dp, "has_default_fparam"): + return bool(self.dp.has_default_fparam()) + if isinstance(self.dp, ModelWrapper): + try: + return self.dp.model["Default"].has_default_fparam() + except AttributeError: + return False + return False
234-245: Static fallbacks for get_intensive/get_var_namePrevent attribute errors in static mode.
Apply:
def get_intensive(self) -> bool: - return self.dp.model["Default"].get_intensive() + if self.static_model and hasattr(self.dp, "get_intensive"): + return bool(self.dp.get_intensive()) + return self.dp.model["Default"].get_intensive() @@ - if hasattr(self.dp.model["Default"], "get_var_name") and callable( - getattr(self.dp.model["Default"], "get_var_name") - ): - return self.dp.model["Default"].get_var_name() + if self.static_model and hasattr(self.dp, "get_var_name"): + return self.dp.get_var_name() + if hasattr(self.dp.model["Default"], "get_var_name") and callable( + self.dp.model["Default"].get_var_name + ): + return self.dp.model["Default"].get_var_name() else: raise NotImplementedError
753-768: get_observed_types: static guardAvoid attribute errors on static models; provide safe fallback.
Apply:
- observed_type_list = self.dp.model["Default"].get_observed_type_list() + if self.static_model and hasattr(self.dp, "get_observed_type_list"): + observed_type_list = self.dp.get_observed_type_list() + else: + observed_type_list = self.dp.model["Default"].get_observed_type_list()
819-835: Descriptor hooks likely unavailable in static graphsStatic (.json) backends may not expose set_eval_descriptor_hook/eval_descriptor; guard or raise gracefully.
Apply:
- model = ( - self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp - ) + model = ( + self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp + ) + if self.static_model and not all( + hasattr(model, name) + for name in ("set_eval_descriptor_hook", "eval_descriptor") + ): + raise NotImplementedError("Descriptor evaluation not available for static (.json) models")
836-890: eval_fitting_last_layer: dynamic‑only; add guardPrevent runtime errors on static backends.
Apply:
- model = self.dp.model["Default"] + if self.static_model: + raise NotImplementedError( + "eval_fitting_last_layer is not available for static (.json) models" + ) + model = self.dp.model["Default"]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pd/infer/deep_eval.py(20 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/infer/deep_eval.py
🧬 Code graph analysis (1)
deepmd/pd/infer/deep_eval.py (3)
deepmd/pd/model/network/network.py (1)
TypeEmbedNetConsistent(127-341)deepmd/utils/econf_embd.py (1)
sort_element_type(273-283)deepmd/utils/model_branch_dict.py (1)
get_model_dict(12-42)
🪛 Ruff (0.13.1)
deepmd/pd/infer/deep_eval.py
112-112: Unused method argument: no_jit
(ARG002)
113-113: Unused method argument: kwargs
(ARG002)
199-199: Do not call setattr with a constant attribute value. It is not any safer than normal property access.
Replace setattr with assignment
(B010)
240-240: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
721-721: Avoid specifying long messages outside the exception class
(TRY003)
🪛 GitHub Check: CodeQL
deepmd/pd/infer/deep_eval.py
[failure] 199-199: Non-callable called
Call to a non-callable of builtin-class bool.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (18)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.12)
🔇 Additional comments (8)
deepmd/pd/infer/deep_eval.py (8)
151-156: Use pathlib; avoid fragile split(".json") and trim type_map stringRobust path handling and safer string parsing.
Apply:
- self.dp = paddle.jit.load(self.model_path.split(".json")[0]) + base = str(Path(self.model_path).with_suffix("")) + self.dp = paddle.jit.load(base) self.rcut = self.dp.get_rcut().item() - self.type_map: list[str] = "".join( - [chr(x) for x in self.dp.get_type_map().numpy()] - ).split(" ") + self.type_map: list[str] = ( + "".join(chr(x) for x in self.dp.get_type_map().numpy()) + .strip() + .split() + ) + # default to empty script for static models; populated if backend exposes an accessor + self.model_def_script = ( + self.dp.get_model_def_script() + if hasattr(self.dp, "get_model_def_script") + else {} + )
157-165: Parameterize GPU config and add CPU fallbackAvoid magic numbers; allow kwargs to override; auto CPU when CUDA unavailable.
Apply:
- config = paddle_inference.Config( - self.model_path, self.model_path.replace(".json", ".pdiparams") - ) + config = paddle_inference.Config( + self.model_path, f"{base}.pdiparams" + ) config.enable_custom_passes( ["add_shadow_output_after_dead_parameter_pass"], True ) - config.enable_use_gpu(4096, 0) + gpu_mem_mb = int(kwargs.get("gpu_mem_mb", 1024)) + device_id = int(kwargs.get("device_id", 0)) + if paddle.device.is_compiled_with_cuda(): + config.enable_use_gpu(gpu_mem_mb, device_id) + else: + config.disable_gpu()
246-265: Static models: use frozen model_output_type()Handle .json path by calling backend’s model_output_type; freeze must include it.
Apply:
@property def model_type(self) -> type["DeepEvalWrapper"]: """The the evaluator of the model type.""" - model_output_type = self.dp.model["Default"].model_output_type() + if self.static_model: + if hasattr(self.dp, "model_output_type"): + model_output_type = self.dp.model_output_type() + else: + raise RuntimeError( + "Static model missing model_output_type(); freeze must include it." + ) + else: + model_output_type = self.dp.model["Default"].model_output_type()
494-496: Static predictor: box shape mismatchFreeze uses [-1, 9]; feed [nframes, 9], not [nframes, 3, 3].
Apply:
- if cells is not None: - self.box_handle.copy_from_cpu(cells.reshape([nframes, 3, 3])) + if cells is not None: + self.box_handle.copy_from_cpu(cells.reshape([nframes, 9]))
654-654: Use Paddle‑compatible NumPy conversion.detach().cpu().numpy() is PyTorch‑style; use helper or .numpy().
Apply:
- out = batch_output[pd_name].reshape(shape).detach().cpu().numpy() + out = to_numpy_array(batch_output[pd_name].reshape(shape))
165-176: Static predictor: add optional do_atomic_virial input handleAllows toggling atomic virial at runtime when exported.
Apply:
self.predictor = paddle_inference.create_predictor(config) self.coord_handle = self.predictor.get_input_handle("coord") self.atype_handle = self.predictor.get_input_handle("atype") self.box_handle = self.predictor.get_input_handle("box") + # Optional boolean control if present + try: + input_names = set(self.predictor.get_input_names()) + if "do_atomic_virial" in input_names: + self.do_atomic_virial_handle = self.predictor.get_input_handle( + "do_atomic_virial" + ) + except Exception: + pass
1-1: Run Ruff lint/format — verification failed: file missingRuff returned E902 (No such file or directory) for deepmd/pd/infer/deep_eval.py in the verification environment; re-run ruff locally (ruff check . && ruff format .) and fix ARG002, B009, B010, TRY003 reported in deepmd/pd/infer/deep_eval.py.
541-555: Static predictor: feed do_atomic_virial when availableYou compute do_atomic_virial but never pass it to the static graph.
Apply:
do_atomic_virial = any( x.category == OutputVariableCategory.DERV_C for x in request_defs ) if self.static_model: + if hasattr(self, "do_atomic_virial_handle"): + self.do_atomic_virial_handle.copy_from_cpu( + np.asarray([do_atomic_virial], dtype="bool") + ) self.predictor.run()
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
source/tests/pd/model/test_descriptor_dpa1.py (1)
371-387: Replace O(n^2) index lookups with an index map and clearer errors.Improves robustness and gives actionable failures if a key is missing.
Apply:
- all_keys = [key for key in target_dict.keys() if "buffer_" not in key] - record = [False for ii in all_keys] - for kk, vv in source_dict.items(): - tk = "se_atten." + kk - record[all_keys.index(tk)] = True - target_dict[tk] = vv + all_keys = [k for k in target_dict.keys() if "buffer_" not in k] + idx = {k: i for i, k in enumerate(all_keys)} + record = [False] * len(all_keys) + + def mark(key: str) -> None: + i = idx.get(key) + assert i is not None, f"Missing expected key in target_dict: {key}" + record[i] = True + + for kk, vv in source_dict.items(): + tk = "se_atten." + kk + mark(tk) + target_dict[tk] = vv @@ - 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("se_atten.compress_data.0")] = True - record[all_keys.index("se_atten.compress_info.0")] = True + for kk in type_embd_dict.keys(): + tk = "type_embedding." + kk + mark(tk) + target_dict[tk] = type_embd_dict[kk] + mark("se_atten.compress_data.0") + mark("se_atten.compress_info.0")source/tests/pd/model/test_descriptor_dpa2.py (1)
193-208: Use an index map to avoid repeated linear searches and improve diagnostics.Keeps tests fast and failure messages precise.
Apply:
- all_keys = [key for key in target_dict.keys() if "buffer_" not in key] - record = [False for ii in all_keys] - for kk, vv in source_dict.items(): - record[all_keys.index(kk)] = True - target_dict[kk] = vv + all_keys = [k for k in target_dict.keys() if "buffer_" not in k] + idx = {k: i for i, k in enumerate(all_keys)} + record = [False] * len(all_keys) + + def mark(key: str) -> None: + i = idx.get(key) + assert i is not None, f"Missing expected key in target_dict: {key}" + record[i] = True + + for kk, vv in source_dict.items(): + mark(kk) + target_dict[kk] = vv @@ - 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 + for kk in type_embd_dict.keys(): + tk = "type_embedding." + kk + mark(tk) + target_dict[tk] = type_embd_dict[kk] + mark("repinit.compress_data.0") + mark("repinit.compress_info.0")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/pd/model/atomic_model/dp_atomic_model.py(3 hunks)source/tests/pd/model/test_descriptor_dpa1.py(1 hunks)source/tests/pd/model/test_descriptor_dpa2.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pd/model/atomic_model/dp_atomic_model.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
source/tests/pd/model/test_descriptor_dpa2.pysource/tests/pd/model/test_descriptor_dpa1.py
🧠 Learnings (1)
📓 Common learnings
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-10-08T15:32:11.479Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. ⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (2)
source/tests/pd/model/test_descriptor_dpa2.py (1)
193-193: Consistent buffer_ filtering for DPA2 — approved; verify translators & lint
- Approved: source/tests/pd/model/test_descriptor_dpa2.py:193 — keeping all_keys = [key for key in target_dict.keys() if "buffer_" not in key] is correct to ignore static buffer keys.
- Verify: the provided rg search returned no matches for translate_.*dpa[12]; confirm the helper translators for DPA1/DPA2 exist and apply the same "buffer_" filtering. The ruff run reported many warnings — run ruff format . and address the lint issues before merging.
source/tests/pd/model/test_descriptor_dpa1.py (1)
371-371: Approve — filtering out "buffer_" keys is correct; verification incompleteFiltering out "buffer_" keys aligns with the static-buffer design and prevents false negatives in the final assert. The supplied verification run produced many ruff warnings and the rg scan failed with "source/deepmd/pd: No such file or directory", so I could not confirm there are no other occurrences of "buffer_" outside model modules. Re-run lint/format and a repository-wide search and attach the output:
#!/bin/bash set -euo pipefail ruff check . ruff format --diff . rg -nP -C1 '\bbuffer_' . | sed -n '1,120p'
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/pd/model/model/ener_model.py (2)
80-94: Force assignment bug: else is paired with do_grad_c instead of do_grad_rWhen do_grad_r("energy") is True and do_grad_c("energy") is False, the current else overwrites force with dforce. Pair the else with do_grad_r.
Apply:
- if self.do_grad_r("energy"): - model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy"): + # force + if self.do_grad_r("energy"): + model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2) + else: + model_predict["force"] = model_ret["dforce"] + # virial + if self.do_grad_c("energy"): model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) if do_atomic_virial: model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze( -3 ) else: model_predict["atom_virial"] = paddle.zeros( [model_predict["energy"].shape[0], 1, 9], dtype=paddle.float64 ) - else: - model_predict["force"] = model_ret["dforce"]
127-142: Extended force assignment bug: else is paired with do_grad_c instead of do_grad_rSame issue in forward_lower: dforce overwrites extended_force when do_grad_r is True and do_grad_c is False.
Apply:
- if self.do_grad_r("energy"): - model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy"): + # force + if self.do_grad_r("energy"): + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + else: + assert model_ret["dforce"] is not None + model_predict["dforce"] = model_ret["dforce"] + # virial + if self.do_grad_c("energy"): model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) if do_atomic_virial: model_predict["extended_virial"] = model_ret[ "energy_derv_c" ].squeeze(-3) else: model_predict["extended_virial"] = paddle.zeros( [model_predict["energy"].shape[0], 1, 9], dtype=paddle.float64 ) - else: - assert model_ret["dforce"] is not None - model_predict["dforce"] = model_ret["dforce"]
🧹 Nitpick comments (17)
deepmd/pd/entrypoints/main.py (1)
356-361: Remove leftover debug printsThese commented prints should be deleted to avoid noise.
Apply:
- # print(model.get_buffer_rcut.__func__.__qualname__) - # print(model.get_buffer_rcut.__func__.__module__) - # print(model.get_buffer_rcut.__func__.__code__.co_filename) - # print(model.get_buffer_rcut.__func__.__code__.co_firstlineno) - # print(model.get_buffer_type_map()) - # exit()deepmd/pd/infer/deep_eval.py (2)
112-115: Wire up no_jit; remove unused parameter warningMake no_jit effective by optionally converting the model to static.
Apply:
- model = get_model(self.input_param).to(DEVICE) - # model = paddle.jit.to_static(model) + model = get_model(self.input_param).to(DEVICE) + if not no_jit: + model = paddle.jit.to_static(model)
876-889: Guard eval_fitting_last_layer for static modelsStatic graphs don’t expose wrapper internals; raise a clear error if called on .json.
Apply:
- model = self.dp.model["Default"] + if self.static_model: + raise NotImplementedError("eval_fitting_last_layer is unsupported for static (.json) models.") + model = self.dp.model["Default"]deepmd/pd/model/descriptor/dpa1.py (2)
300-305: Handle None type_map when registering buffer_type_mapAvoid attribute errors later by always registering the buffer (empty tensor if None).
Apply:
- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + ) + else: + self.register_buffer("buffer_type_map", paddle.empty([0], dtype="int32"))
328-339: Type annotation mismatch: get_buffer_rcut/_smth return Tensor, not floatFix return annotations to paddle.Tensor for consistency.
Apply:
- def get_buffer_rcut(self) -> float: + def get_buffer_rcut(self) -> paddle.Tensor: @@ - def get_buffer_rcut_smth(self) -> float: + def get_buffer_rcut_smth(self) -> paddle.Tensor:deepmd/pd/model/descriptor/dpa2.py (2)
268-273: Handle None type_map when registering buffer_type_mapMirror dpa1: always register, use empty tensor when type_map is None.
Apply:
- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + ) + else: + self.register_buffer("buffer_type_map", paddle.empty([0], dtype="int32"))
787-787: Minor: Prefer getters for consistency (optional)Using self.repinit.get_rcut()/get_nsel() keeps the pattern consistent with other calls.
deepmd/pd/model/descriptor/repflows.py (2)
383-386: Return type mismatch inget_ntypes.Docstring says int, but static branch returns a Tensor. Either cast to int here or adjust the docstring and add a dedicated
get_buffer_ntypes().
219-231: Avoid mutable default argumentexclude_types=[].Use
Noneas default and normalize to[]inside to prevent shared state bugs.deepmd/pd/model/descriptor/se_t_tebd.py (2)
649-652: Return type mismatch inget_ntypes.Docstring says int, but static branch returns Tensor.
536-544: Avoid mutable default argumentexclude_types=[].Same recommendation as other descriptors.
deepmd/pd/model/descriptor/dpa3.py (2)
254-255: Return type mismatch inget_ntypes.Docstring says int, but static branch returns Tensor.
114-122: Avoid mutable default argumentexclude_types=[].Use
Nonedefault and normalize internally.deepmd/pd/model/descriptor/se_a.py (4)
146-149: Guardget_buffer_type_mapwhen buffer may be empty.Optional: if you don’t register the buffer, this accessor will fail.
546-549: Return type mismatch inget_ntypes.Docstring says int; static branch returns Tensor.
432-440: Avoid mutable default argumentexclude_types=[].Same pattern; switch to
Nonedefault.
530-537: Approve — buffer accessors are correct; address linter warningsget_buffer_rcut and get_buffer_rcut_smth match the newly registered buffers and are compatible with static inference.
Ruff run returned many lint warnings across the repo; notable issues in deepmd/pd/model/descriptor/se_a.py — B006 (mutable default args) at ~77,83,236,426,432; ARG002 (unused args) at 243,272; RUF059 (unused unpacked g2/h2) at 308; F841 (env_mat assigned but unused) at 373. Fix the reported lints and re-run:
ruff format . && ruff check .
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (16)
deepmd/pd/entrypoints/main.py(5 hunks)deepmd/pd/infer/deep_eval.py(20 hunks)deepmd/pd/model/atomic_model/base_atomic_model.py(3 hunks)deepmd/pd/model/atomic_model/dp_atomic_model.py(3 hunks)deepmd/pd/model/descriptor/dpa1.py(3 hunks)deepmd/pd/model/descriptor/dpa2.py(5 hunks)deepmd/pd/model/descriptor/dpa3.py(4 hunks)deepmd/pd/model/descriptor/repflows.py(3 hunks)deepmd/pd/model/descriptor/repformers.py(3 hunks)deepmd/pd/model/descriptor/se_a.py(6 hunks)deepmd/pd/model/descriptor/se_atten.py(4 hunks)deepmd/pd/model/descriptor/se_t_tebd.py(6 hunks)deepmd/pd/model/model/ener_model.py(1 hunks)deepmd/pd/model/model/make_model.py(1 hunks)deepmd/pd/model/model/model.py(1 hunks)deepmd/pd/model/task/fitting.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (4)
- deepmd/pd/model/atomic_model/base_atomic_model.py
- deepmd/pd/model/atomic_model/dp_atomic_model.py
- deepmd/pd/model/task/fitting.py
- deepmd/pd/model/model/model.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/model/model/ener_model.pydeepmd/pd/model/model/make_model.pydeepmd/pd/entrypoints/main.pydeepmd/pd/model/descriptor/repflows.pydeepmd/pd/model/descriptor/dpa2.pydeepmd/pd/model/descriptor/se_atten.pydeepmd/pd/model/descriptor/dpa3.pydeepmd/pd/infer/deep_eval.pydeepmd/pd/model/descriptor/dpa1.pydeepmd/pd/model/descriptor/repformers.pydeepmd/pd/model/descriptor/se_t_tebd.pydeepmd/pd/model/descriptor/se_a.py
🧬 Code graph analysis (11)
deepmd/pd/model/model/ener_model.py (8)
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-139)deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_buffer_type_map(160-162)deepmd/pd/model/descriptor/dpa1.py (1)
get_buffer_type_map(356-358)deepmd/pd/model/descriptor/dpa3.py (1)
get_buffer_type_map(260-262)deepmd/pd/model/descriptor/se_a.py (1)
get_buffer_type_map(146-148)deepmd/pd/model/descriptor/se_t_tebd.py (1)
get_buffer_type_map(216-218)deepmd/pd/model/model/make_model.py (1)
get_buffer_type_map(560-562)deepmd/pd/model/task/fitting.py (1)
get_buffer_type_map(478-480)
deepmd/pd/model/model/make_model.py (2)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_rcut(164-166)get_buffer_type_map(160-162)deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-139)
deepmd/pd/model/descriptor/repflows.py (4)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_rcut(164-166)get_buffer_sel(168-170)deepmd/pd/model/descriptor/dpa3.py (3)
get_buffer_rcut(232-234)get_buffer_sel(248-250)get_ntypes(252-254)deepmd/pd/model/descriptor/repformers.py (2)
get_buffer_rcut(332-334)get_ntypes(348-350)deepmd/pd/model/model/make_model.py (1)
get_buffer_rcut(556-558)
deepmd/pd/model/descriptor/dpa2.py (7)
deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_buffer_rcut(164-166)deepmd/pd/model/descriptor/repflows.py (1)
get_buffer_rcut(363-365)deepmd/pd/model/descriptor/repformers.py (2)
get_buffer_rcut(332-334)get_buffer_rcut_smth(336-338)deepmd/pd/model/descriptor/se_a.py (2)
get_buffer_rcut(530-532)get_buffer_rcut_smth(534-536)deepmd/pd/model/descriptor/se_atten.py (2)
get_buffer_rcut(287-289)get_buffer_rcut_smth(291-293)deepmd/pd/model/descriptor/se_t_tebd.py (2)
get_buffer_rcut(633-635)get_buffer_rcut_smth(637-639)deepmd/pd/utils/nlist.py (1)
get_multiple_nlist_key(332-336)
deepmd/pd/model/descriptor/se_atten.py (9)
deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_buffer_rcut(164-166)deepmd/pd/model/descriptor/dpa1.py (2)
get_buffer_rcut(328-330)get_buffer_rcut_smth(336-338)deepmd/pd/model/descriptor/dpa2.py (2)
get_buffer_rcut(343-345)get_buffer_rcut_smth(347-349)deepmd/pd/model/descriptor/dpa3.py (2)
get_buffer_rcut(232-234)get_buffer_rcut_smth(236-238)deepmd/pd/model/descriptor/repflows.py (1)
get_buffer_rcut(363-365)deepmd/pd/model/descriptor/repformers.py (2)
get_buffer_rcut(332-334)get_buffer_rcut_smth(336-338)deepmd/pd/model/descriptor/se_a.py (2)
get_buffer_rcut(530-532)get_buffer_rcut_smth(534-536)deepmd/pd/model/descriptor/se_t_tebd.py (2)
get_buffer_rcut(633-635)get_buffer_rcut_smth(637-639)deepmd/pd/model/model/make_model.py (1)
get_buffer_rcut(556-558)
deepmd/pd/model/descriptor/dpa3.py (4)
deepmd/pd/model/atomic_model/dp_atomic_model.py (4)
get_buffer_rcut(164-166)get_buffer_sel(168-170)get_sel(156-158)get_buffer_type_map(160-162)deepmd/pd/model/descriptor/repflows.py (4)
get_buffer_rcut(363-365)get_buffer_sel(379-381)get_sel(375-377)get_ntypes(383-385)deepmd/pd/model/descriptor/repformers.py (4)
get_buffer_rcut(332-334)get_buffer_rcut_smth(336-338)get_sel(344-346)get_ntypes(348-350)source/api_cc/src/DeepPotPD.cc (2)
get_type_map(658-667)get_type_map(658-658)
deepmd/pd/infer/deep_eval.py (6)
deepmd/jax/infer/deep_eval.py (5)
DeepEval(60-434)get_model(426-434)eval(189-252)model_type(152-166)_get_output_shape(398-420)deepmd/dpmodel/infer/deep_eval.py (5)
DeepEval(57-413)get_model(405-413)eval(168-231)model_type(131-145)_get_output_shape(375-399)deepmd/pt/infer/deep_eval.py (9)
DeepEval(85-841)get_model(723-731)eval(306-380)get_var_name(232-239)model_type(242-260)_get_output_shape(623-647)_eval_model_spin(537-621)get_observed_types(707-721)eval_fitting_last_layer(788-841)deepmd/pd/model/network/network.py (1)
TypeEmbedNetConsistent(127-341)deepmd/utils/econf_embd.py (1)
sort_element_type(273-283)deepmd/utils/model_branch_dict.py (1)
get_model_dict(12-42)
deepmd/pd/model/descriptor/dpa1.py (4)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_rcut(164-166)get_buffer_type_map(160-162)deepmd/pd/model/descriptor/se_atten.py (3)
get_buffer_rcut(287-289)get_rcut_smth(283-285)get_buffer_rcut_smth(291-293)deepmd/pd/model/model/make_model.py (2)
get_buffer_rcut(556-558)get_buffer_type_map(560-562)deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-139)
deepmd/pd/model/descriptor/repformers.py (9)
deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_buffer_rcut(164-166)deepmd/pd/model/descriptor/dpa1.py (2)
get_buffer_rcut(328-330)get_buffer_rcut_smth(336-338)deepmd/pd/model/descriptor/dpa2.py (2)
get_buffer_rcut(343-345)get_buffer_rcut_smth(347-349)deepmd/pd/model/descriptor/dpa3.py (2)
get_buffer_rcut(232-234)get_buffer_rcut_smth(236-238)deepmd/pd/model/descriptor/repflows.py (1)
get_buffer_rcut(363-365)deepmd/pd/model/descriptor/se_a.py (2)
get_buffer_rcut(530-532)get_buffer_rcut_smth(534-536)deepmd/pd/model/descriptor/se_atten.py (2)
get_buffer_rcut(287-289)get_buffer_rcut_smth(291-293)deepmd/pd/model/descriptor/se_t_tebd.py (2)
get_buffer_rcut(633-635)get_buffer_rcut_smth(637-639)deepmd/pd/model/model/make_model.py (1)
get_buffer_rcut(556-558)
deepmd/pd/model/descriptor/se_t_tebd.py (4)
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-139)deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_type_map(160-162)get_buffer_rcut(164-166)deepmd/pd/model/model/ener_model.py (1)
get_buffer_type_map(37-39)deepmd/pd/model/model/make_model.py (2)
get_buffer_type_map(560-562)get_buffer_rcut(556-558)
deepmd/pd/model/descriptor/se_a.py (6)
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-139)deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_type_map(160-162)get_buffer_rcut(164-166)deepmd/pd/model/descriptor/dpa1.py (3)
get_buffer_type_map(356-358)get_buffer_rcut(328-330)get_buffer_rcut_smth(336-338)deepmd/pd/model/model/make_model.py (2)
get_buffer_type_map(560-562)get_buffer_rcut(556-558)deepmd/pd/model/descriptor/repflows.py (1)
get_buffer_rcut(363-365)deepmd/pd/model/descriptor/se_atten.py (2)
get_buffer_rcut(287-289)get_buffer_rcut_smth(291-293)
🪛 Ruff (0.13.1)
deepmd/pd/infer/deep_eval.py
112-112: Unused method argument: no_jit
(ARG002)
113-113: Unused method argument: kwargs
(ARG002)
199-199: Do not call setattr with a constant attribute value. It is not any safer than normal property access.
Replace setattr with assignment
(B010)
240-240: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
721-721: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.9)
🔇 Additional comments (19)
deepmd/pd/model/model/ener_model.py (3)
37-40: Exposing buffer_type_map via EnergyModel looks goodDelegating to the parent is consistent with the API surface across models/descriptors.
49-55: Potential no-op squeeze() on output defs—verify intentCalling .squeeze(-2/-3) without assignment likely does nothing unless these objects mutate in place. If they’re OutputVariableDef-like, confirm squeeze() mutates the def; otherwise, reassign.
1-1: Reminder: run ruff before commitPer guidelines, please run “ruff check .” and “ruff format .” for Python files.
deepmd/pd/infer/deep_eval.py (7)
151-166: Use robust path handling and avoid magic GPU numbers
- Path: split(".json")[0] is fragile. Prefer pathlib.Path.with_suffix("").
- GPU: 4096, 0 are magic numbers; make configurable or document.
Suggested change (requires adding
from pathlib import Pathat imports):- self.dp = paddle.jit.load(self.model_path.split(".json")[0]) + self.dp = paddle.jit.load(str(Path(self.model_path).with_suffix("")))Optionally make GPU memory/device configurable via kwargs.
541-555: Static atom_virial output may be absentIf the frozen graph was exported with do_atomic_virial=False, atom_virial isn’t an output; accessing its handle will fail. Consider guarding by checking predictor.get_output_names() or failing fast when atomic outputs are requested on a static model without support.
193-200: Fix has_spin detection; avoid calling non-callables and setattr misuseCurrent code may call a non-callable and trips Ruff/CodeQL. Simplify and coerce to bool.
Apply:
- self._has_spin = ( - getattr(self.dp.model["Default"], "has_spin", False) - if isinstance(self.dp, ModelWrapper) - else False - ) - if callable(self._has_spin) and not isinstance(self._has_spin, bool): - setattr(self, "_has_spin", self._has_spin()) + has_spin_attr = ( + getattr(self.dp.model["Default"], "has_spin", None) + if isinstance(self.dp, ModelWrapper) + else None + ) + if callable(has_spin_attr): + self._has_spin = bool(has_spin_attr()) + else: + self._has_spin = bool(has_spin_attr) if has_spin_attr is not None else False
246-265: Support static models in model_type selectionStatic .json lacks model["Default"]; call frozen model_output_type() when available.
Apply:
- model_output_type = self.dp.model["Default"].model_output_type() + if self.static_model: + if hasattr(self.dp, "model_output_type"): + model_output_type = self.dp.model_output_type() + else: + raise RuntimeError("Static model missing model_output_type(); freeze must include it.") + else: + model_output_type = self.dp.model["Default"].model_output_type()
485-506: Static predictor: box shape must be [nframes, 9]Freeze uses InputSpec([-1, 9]) for box; here it’s [nframes, 3, 3].
Apply:
- if cells is not None: - self.box_handle.copy_from_cpu(cells.reshape([nframes, 3, 3])) + if cells is not None: + self.box_handle.copy_from_cpu(cells.reshape([nframes, 9]))
583-594: Block unsupported spin path for static modelsStatic (.json) spin inference isn’t implemented; fail fast.
Apply:
def _eval_model_spin( @@ - model = self.dp.to(DEVICE) + if self.static_model: + raise NotImplementedError( + "Spin inference for static (.json) models is not supported; use a .pd model." + ) + model = self.dp.to(DEVICE)
654-666: Paddle Tensor to NumPy: replace .detach().cpu().numpy()Use to_numpy_array() for Paddle, not PyTorch-style APIs.
Apply:
- out = batch_output[pd_name].reshape(shape).detach().cpu().numpy() + out = to_numpy_array(batch_output[pd_name].reshape(shape))deepmd/pd/model/descriptor/repformers.py (1)
332-339: LGTM: buffer getters align with static inference patternReturning registered buffers for rcut/rcut_smth is consistent with other descriptors.
deepmd/pd/model/model/make_model.py (1)
556-563: LGTM: CM exposes buffer_rcut and buffer_type_mapThese delegations complete the buffer exposure at the model level.
deepmd/pd/model/descriptor/se_atten.py (2)
187-193: LGTM: buffer_ntypes registration is consistent across descriptorsPattern matches other modules; enables static retrieval of ntypes.
287-294: LGTM: buffer getters for rcut/rcut_smthMatches the buffer exposure approach throughout the PR.
deepmd/pd/entrypoints/main.py (1)
419-438: Static-ize model_output_type for static modelsdeep_eval.py needs model_output_type() for evaluator selection in static mode. Add it to the to_static wrapping loop.
Apply:
for method_name in [ "get_buffer_rcut", "get_buffer_type_map", "get_buffer_dim_fparam", "get_buffer_dim_aparam", "get_buffer_intensive", "get_buffer_sel_type", "get_buffer_numb_dos", "get_buffer_task_dim", + "model_output_type", ]:deepmd/pd/model/descriptor/se_t_tebd.py (1)
212-215: Reconstructtype_mapfrom buffer in non-dynamic mode.Static mode needs
list[str]rather than raw char codes.Apply this diff:
def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" - return self.type_map + if paddle.in_dynamic_mode(): + return self.type_map if self.type_map is not None else [] + # static path: decode from buffer + if self.buffer_type_map.numel() == 0: + return [] + codes = self.buffer_type_map.numpy().tolist() + return "".join(chr(x) for x in codes).split(" ")deepmd/pd/model/descriptor/dpa3.py (3)
256-259: Reconstructtype_mapfrom buffer in non-dynamic mode.Mirror the decoding pattern to avoid
Nonein static runs.Apply this diff:
def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" - return self.type_map + if paddle.in_dynamic_mode(): + return self.type_map if self.type_map is not None else [] + if self.buffer_type_map.numel() == 0: + return [] + codes = self.buffer_type_map.numpy().tolist() + return "".join(chr(x) for x in codes).split(" ")
179-184: Handletype_map=Noneby registering an empty buffer.Prevents AttributeError for
get_buffer_type_mapin static graphs.Apply this diff:
- if type_map is not None: + if type_map is not None: self.register_buffer( "buffer_type_map", paddle.to_tensor([ord(c) for c in " ".join(self.type_map)]), ) + else: + self.register_buffer( + "buffer_type_map", paddle.to_tensor([], dtype="int32") + )
236-239: Incorrect — repflows.get_buffer_rcut_smth() exists; no change required.Implemented in deepmd/pd/model/descriptor/repflows.py (get_buffer_rcut_smth returns self.buffer_rcut_smth).
Likely an incorrect or invalid review comment.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
source/tests/pd/model/test_deeppot.py (1)
132-139: Assert CPU placement in test — deepmd/pd/infer/deep_eval.py uses DEVICE for model.to(...) and paddle.to_tensor(..., place=DEVICE) and no device_guard/set_device was found; the existing mocks (deepmd.pd.infer.deep_eval.DEVICE and deepmd.pd.utils.env.DEVICE → paddle.CPUPlace()) should be effective, but add explicit assertions in test_dp_test_cpu that the model parameters and created tensors are placed on CPU (inspect .place or str(place)).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pd/model/test_deeppot.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
source/tests/pd/model/test_deeppot.py
🧠 Learnings (3)
📓 Common learnings
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-10-08T15:32:11.479Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. 📚 Learning: 2024-10-08T15:32:11.479Z
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-10-08T15:32:11.479Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Applied to files:
source/tests/pd/model/test_deeppot.py
📚 Learning: 2025-09-18T11:37:10.520Z
Learnt from: CR PR: deepmodeling/deepmd-kit#0 File: AGENTS.md:0-0 Timestamp: 2025-09-18T11:37:10.520Z Learning: Applies to source/tests/tf/test_dp_test.py : Keep the core TensorFlow test `source/tests/tf/test_dp_test.py` passing; use it for quick validation Applied to files:
source/tests/pd/model/test_deeppot.py
🧬 Code graph analysis (1)
source/tests/pd/model/test_deeppot.py (2)
deepmd/pd/entrypoints/main.py (3)
main(565-613)freeze(341-448)get_trainer(82-191)deepmd/pd/infer/deep_eval.py (7)
eval(311-385)get_type_map(210-212)get_ntypes(206-208)get_dim_fparam(214-218)get_dim_aparam(220-224)model_type(247-265)eval_typeebd(695-723)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
🔇 Additional comments (7)
source/tests/pd/model/test_deeppot.py (7)
1-23: LGTM! Clean imports and setup.The imports are well-organized and appropriate for a comprehensive DeepPot test module. The licensing header is correct.
25-49: LGTM! Solid test setup with proper configuration.The setUp method properly:
- Loads configuration from a reference file
- Reduces training steps to 1 for fast testing
- Sets up training/validation data paths
- Runs the trainer to generate a model
- Prepares input data for evaluation
The approach of using minimal training steps is appropriate for unit testing.
50-54: LGTM! Proper cleanup in tearDown.The tearDown method correctly removes test artifacts to keep the test environment clean.
55-104: LGTM! Comprehensive DeepPot evaluation test.The test method properly:
- Creates a DeepPot instance from the trained model
- Sets up realistic coordinate, cell, and atom type data
- Performs evaluation with atomic=True to get detailed outputs
- Validates output shapes for energy (e), forces (f), virial (v), atomic energy (ae), and atomic virial (av)
- Tests accessor methods for type map, number of types, and parameter dimensions
- Verifies the model type is correctly identified as DeepPot
The test data appears to represent a water system with realistic coordinates and cell parameters.
105-109: LGTM! Good integration test for unified interface.This test validates that the unified DeepPotUni interface correctly wraps and returns a DeepPot instance, ensuring backward compatibility.
110-117: LGTM! Appropriate type embedding test.The test correctly:
- Validates the type embedding output shape matches expected dimensions (ntypes + 1, embedding_size)
- Verifies the last row (padding/unknown type) is zeros
- Uses the configuration's type_map length for validation
This aligns with the PR's addition of type embedding evaluation functionality.
119-131: Good frozen model test setup with appropriate skip.The skip decorator with clear reasoning is appropriate since frozen models (.json) have limited attribute access capabilities. The setUp properly inherits from the parent class and freezes the model with atomic virial support enabled.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
.github/workflows/test_cuda.yml (2)
67-67: Guard against undefined CUDA_PATH in LD_LIBRARY_PATH export.Some CUDA images don’t set CUDA_PATH. Default to /usr/local/cuda to avoid an empty path segment.
Apply this diff:
- export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH + export CUDA_PATH=${CUDA_PATH:-/usr/local/cuda} + export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
50-50: Use uv_with_retry wrapper for paddlepaddle-gpu install (cp311 wheel verified)paddlepaddle-gpu==3.2.0 cp311 wheel is present on the cu126 index and downloads successfully; switch this install to the repo's retry wrapper and use --system to match other installs and improve resilience.
- pip install "paddlepaddle-gpu==3.2.0" -i https://www.paddlepaddle.org.cn/packages/stable/cu126/ + source/install/uv_with_retry.sh pip install --system "paddlepaddle-gpu==3.2.0" -i https://www.paddlepaddle.org.cn/packages/stable/cu126/
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
.github/workflows/test_cuda.yml(1 hunks).github/workflows/test_python.yml(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- .github/workflows/test_python.yml
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (5)
source/tests/pd/model/test_deeppot.py (5)
25-51: Clean up temporary files created during training.The test creates several files during training (model checkpoints, logs) but only removes
lcurve.outand the JSON config intearDown. This could leave behind model files and other artifacts between test runs.Consider using a temporary directory for all test artifacts:
+import tempfile +import shutil class TestDeepPot(unittest.TestCase): def setUp(self) -> None: + self.temp_dir = tempfile.mkdtemp() + self.original_cwd = os.getcwd() + os.chdir(self.temp_dir) + input_json = str(Path(__file__).parent / "water/se_atten.json")def tearDown(self) -> None: - for f in os.listdir("."): - if f in ["lcurve.out", self.input_json]: - os.remove(f) + os.chdir(self.original_cwd) + shutil.rmtree(self.temp_dir)
57-106: Verify that the evaluation method returns expected keys.The test accesses return values by index (
ret[0],ret[1], etc.) which assumes a specific order and number of outputs. This is brittle if the evaluation method changes its return structure.Access outputs by key names for better maintainability:
ret = dp.eval(coord, cell, atype, atomic=True) - e, f, v, ae, av = ret[0], ret[1], ret[2], ret[3], ret[4] + self.assertIn("energy", ret) + self.assertIn("force", ret) + self.assertIn("virial", ret) + self.assertIn("atom_energy", ret) + self.assertIn("atom_virial", ret) + + e, f, v = ret["energy"], ret["force"], ret["virial"] + ae, av = ret["atom_energy"], ret["atom_virial"]
93-94: Validate coordinate and cell data types for consistency.The test uses hardcoded float64 values but doesn't ensure the arrays have the correct dtype, which could cause issues if the model expects specific precision.
Explicitly set dtypes to match expected model inputs:
- ret = dp.eval(coord, cell, atype, atomic=True) + ret = dp.eval( + coord.astype(np.float64), + cell.astype(np.float64), + atype.astype(np.int32), + atomic=True + )
112-118: Add edge case validation for eval_typeebd output.The test only checks the shape and that the last row is zeros, but doesn't validate the actual embedding values or verify the concatenation behavior mentioned in the docstring.
Add more comprehensive validation:
eval_typeebd = dp.eval_typeebd() self.assertEqual( eval_typeebd.shape, (len(self.config["model"]["type_map"]) + 1, 8) ) + # Verify embedding values are non-zero for actual types + for i in range(len(self.config["model"]["type_map"])): + self.assertFalse(np.allclose(eval_typeebd[i], 0), + f"Type {i} embedding should not be all zeros") + # Verify padding row is zeros np.testing.assert_allclose(eval_typeebd[-1], np.zeros_like(eval_typeebd[-1]))
121-132: Consider removing the skip decorator and implementing proper frozen model testing.The test class is entirely skipped because frozen models don't support getting attributes. This reduces test coverage for an important feature (JSON inference).
Either implement tests that work with frozen models or add specific skip decorators for individual methods:
-@unittest.skip(reason="Freezed model(.json) do not support getting attributes") class TestDeepPotFrozen(TestDeepPot): def setUp(self) -> None: super().setUp() frozen_model = "frozen_model.json" freeze( model=self.model, output=frozen_model, head=None, do_atomic_virial=True, ) self.model = frozen_model + @unittest.skip("Frozen models don't support attribute access") + def test_eval_typeebd(self) -> None: + pass + + def test_dp_test(self) -> None: + # Test only evaluation, skip attribute checks + dp = DeepPot(str(self.model)) + # ... evaluation tests only
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pd/utils/decomp.py(1 hunks)source/tests/pd/model/test_deeppot.py(1 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/utils/decomp.pysource/tests/pd/model/test_deeppot.py
🧠 Learnings (3)
📓 Common learnings
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-10-08T15:32:11.479Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. 📚 Learning: 2024-09-19T04:25:12.408Z
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Applied to files:
source/tests/pd/model/test_deeppot.py
📚 Learning: 2025-09-18T11:37:10.532Z
Learnt from: CR PR: deepmodeling/deepmd-kit#0 File: AGENTS.md:0-0 Timestamp: 2025-09-18T11:37:10.532Z Learning: Applies to source/tests/tf/test_dp_test.py : Keep the core TensorFlow test `source/tests/tf/test_dp_test.py` passing; use it for quick validation Applied to files:
source/tests/pd/model/test_deeppot.py
🧬 Code graph analysis (1)
source/tests/pd/model/test_deeppot.py (2)
deepmd/pd/entrypoints/main.py (3)
main(565-613)freeze(341-448)get_trainer(82-191)deepmd/pd/infer/deep_eval.py (7)
eval(311-385)get_type_map(210-212)get_ntypes(206-208)get_dim_fparam(214-218)get_dim_aparam(220-224)model_type(247-265)eval_typeebd(695-723)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
🔇 Additional comments (1)
deepmd/pd/utils/decomp.py (1)
115-115: LGTM - Simplified paddle.nonzero usage.The change from
paddle.nonzero(mask, as_tuple=True)followed by coordinate extraction and concatenation to directly usingpaddle.nonzero(mask, as_tuple=False)is a good simplification that achieves the same result with cleaner code.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@ ## devel #4987 +/- ## ========================================== + Coverage 84.21% 84.27% +0.06% ========================================== Files 705 705 Lines 69315 69583 +268 Branches 3575 3576 +1 ========================================== + Hits 58372 58640 +268 + Misses 9803 9802 -1 - Partials 1140 1141 +1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
deepmd/pd/utils/stat.py (2)
85-92: Bug: Out-of-scope variable dd used; concat guard is incorrectLine 90 references stat_data[dd] inside a loop over sys_stat. dd is not defined in this scope and this condition may misbehave. Check the first element of sys_stat[key] instead.
- for key in sys_stat: + for key in sys_stat: if isinstance(sys_stat[key], np.float32): pass elif sys_stat[key] is None or sys_stat[key][0] is None: sys_stat[key] = None - elif isinstance(stat_data[dd], paddle.Tensor): + elif isinstance(sys_stat[key][0], paddle.Tensor): sys_stat[key] = paddle.concat(sys_stat[key], axis=0)
251-262: compute_output_stats return annotation and model_forward type are incorrect
- The function returns a tuple of dicts (bias, std), not a single dict.
- model_forward returns dict[str, paddle.Tensor], not a single tensor.
def compute_output_stats( merged: Union[Callable[[], list[dict]], list[dict]], ntypes: int, keys: Union[str, list[str]] = ["energy"], stat_file_path: Optional[DPPath] = None, rcond: Optional[float] = None, preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, - model_forward: Optional[Callable[..., paddle.Tensor]] = None, + model_forward: Optional[Callable[..., dict[str, paddle.Tensor]]] = None, stats_distinguish_types: bool = True, intensive: bool = False, -) -> dict[str, Any]: +) -> tuple[dict[str, paddle.Tensor], dict[str, paddle.Tensor]]:deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
92-99: buffer_ntypes dtype conflicts with BaseAtomicModelBase registers int64; here it's int32. This mismatch can break consumers expecting a consistent dtype.
Apply:
- self.register_buffer( - "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int32") - ) + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + )Alternatively, drop re-registration and rely on BaseAtomicModel’s buffer.
🧹 Nitpick comments (13)
deepmd/pd/utils/stat.py (1)
536-536: Avoid nested rmse; define once at module levelDefine rmse once at module scope and reuse to reduce duplication and improve testability.
Add at module level (outside functions):
def rmse(x: np.ndarray) -> float: return np.sqrt(np.mean(np.square(x)))Then remove the nested rmse in compute_output_stats_global.
source/tests/pd/model/test_atomic_model_atomic_stat.py (1)
246-248: Deduplicate raise_error() helperDefined twice across classes. Consider a single top-level helper to DRY.
Apply this diff to define it once at module scope and reuse:
+def raise_error() -> NoReturn: + raise RuntimeError @@ - def raise_error() -> NoReturn: - raise RuntimeError + # reuse module-level raise_error() @@ - def raise_error() -> NoReturn: - raise RuntimeError + # reuse module-level raise_error()Also applies to: 407-409
deepmd/pd/model/descriptor/se_a.py (1)
471-474: Unify dtype for buffer_ntypesHere it's int64; other places (e.g., DPAtomicModel) use int32. Pick one consistently (recommend int64 for indices/counts) to avoid surprises.
Would you like me to sweep the repo and align dtypes?
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
162-165: Fix docstring: returns type map buffer, not neighbor selectionAvoids confusion in public API docs.
- def get_buffer_type_map(self) -> paddle.Tensor: - """Get the neighbor selection.""" + def get_buffer_type_map(self) -> paddle.Tensor: + """Get the type-map buffer.""" return self.buffer_type_map
76-84: Potential double registration of buffer_type_mapBoth BaseAtomicModel and DPAtomicModel register buffer_type_map; redundant and risks dtype drift. Prefer registering once in Base.
- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor( - _string_to_array(" ".join(self.type_map)), dtype="int32" - ), - ) - self.buffer_type_map.name = "buffer_type_map" + # buffer_type_map is registered in BaseAtomicModeldeepmd/pd/model/atomic_model/base_atomic_model.py (1)
293-294: Remove stray commentLeftover “# raise” is confusing.
- # raise return ret_dictdeepmd/pd/infer/deep_eval.py (7)
143-145: Wire up no_jit to optionally to_staticno_jit is currently unused. Use it to control paddle.jit.to_static().
- # model = paddle.jit.to_static(model) + if not no_jit: + model = paddle.jit.to_static(model)
152-156: Robust path handling for .json modelssplit(".json")[0] is fragile. Use pathlib to drop suffix safely.
- self.dp = paddle.jit.load(self.model_path.split(".json")[0]) + self.dp = paddle.jit.load(str(Path(self.model_path).with_suffix("")))Add the import near the top:
from pathlib import Path
164-164: Make GPU config configurable; avoid magic numbersHardcoded (4096, 0) isn’t portable. Accept via kwargs/env, with sensible default.
For example:
- config.enable_use_gpu(4096, 0) + gpu_mem_mb = int(kwargs.get("gpu_mem_mb", 4096)) + gpu_id = int(kwargs.get("gpu_id", 0)) + config.enable_use_gpu(gpu_mem_mb, gpu_id)
247-267: Static model_type: try frozen model_output_type if availableEnables non-energy static backends when exported properly.
- if self.static_model: - return DeepPot - model_output_type = self.dp.model["Default"].model_output_type() + if self.static_model: + if hasattr(self.dp, "model_output_type"): + model_output_type = self.dp.model_output_type() + else: + return DeepPot + else: + model_output_type = self.dp.model["Default"].model_output_type()
657-657: Paddle tensor to NumPy: avoid .cpu(); use helperPrefer to_numpy_array for Paddle.
- out = batch_output[pd_name].reshape(shape).detach().cpu().numpy() + out = to_numpy_array(batch_output[pd_name].reshape(shape))
822-838: Guard descriptor evaluation for static modelsStatic graphs don’t expose hooks; raise clearly.
- model = ( - self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp - ) + if self.static_model: + raise NotImplementedError("eval_descriptor is not supported for static (.json) models.") + model = self.dp.model["Default"]
879-893: Guard eval_fitting_last_layer for static modelsSame reason as above.
- model = self.dp.model["Default"] + if self.static_model: + raise NotImplementedError("eval_fitting_last_layer is not supported for static (.json) models.") + model = self.dp.model["Default"]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
deepmd/pd/entrypoints/main.py(7 hunks)deepmd/pd/infer/deep_eval.py(20 hunks)deepmd/pd/model/atomic_model/base_atomic_model.py(13 hunks)deepmd/pd/model/atomic_model/dp_atomic_model.py(10 hunks)deepmd/pd/model/descriptor/se_a.py(7 hunks)deepmd/pd/model/model/make_model.py(2 hunks)deepmd/pd/utils/stat.py(10 hunks)source/tests/pd/model/test_atomic_model_atomic_stat.py(7 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pd/model/model/make_model.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/entrypoints/main.pysource/tests/pd/model/test_atomic_model_atomic_stat.pydeepmd/pd/model/descriptor/se_a.pydeepmd/pd/model/atomic_model/dp_atomic_model.pydeepmd/pd/infer/deep_eval.pydeepmd/pd/model/atomic_model/base_atomic_model.pydeepmd/pd/utils/stat.py
🧬 Code graph analysis (7)
deepmd/pd/entrypoints/main.py (3)
deepmd/pd/train/training.py (1)
Trainer(85-1219)deepmd/pd/utils/dataloader.py (1)
DpLoaderSet(66-271)deepmd/utils/data_system.py (1)
process_systems(787-814)
source/tests/pd/model/test_atomic_model_atomic_stat.py (3)
source/tests/pd/test_training.py (11)
tearDown(132-139)tearDown(155-156)tearDown(173-174)tearDown(194-196)tearDown(211-212)tearDown(227-228)setUp(143-153)setUp(160-171)setUp(180-192)setUp(200-209)setUp(216-225)source/tests/pd/test_multitask.py (8)
tearDown(191-198)tearDown(236-237)tearDown(276-277)tearDown(315-316)tearDown(354-355)tearDown(393-394)tearDown(435-436)setUp(202-234)source/tests/pd/model/test_atomic_model_global_stat.py (3)
tearDown(140-141)raise_error(246-247)raise_error(358-359)
deepmd/pd/model/descriptor/se_a.py (3)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_rcut(166-168)get_buffer_type_map(162-164)deepmd/pd/model/model/make_model.py (2)
get_buffer_rcut(564-566)get_buffer_type_map(568-570)deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-139)
deepmd/pd/model/atomic_model/dp_atomic_model.py (6)
deepmd/pd/model/atomic_model/base_atomic_model.py (4)
get_buffer_type_map(137-139)deserialize(363-380)compute_or_load_out_stat(409-434)has_default_fparam(149-151)deepmd/pd/model/descriptor/se_a.py (4)
get_buffer_type_map(154-156)get_buffer_rcut(130-132)get_buffer_rcut(538-540)deserialize(374-390)deepmd/pd/model/model/make_model.py (8)
get_buffer_type_map(568-570)get_buffer_rcut(564-566)set_case_embd(521-522)deserialize(518-519)get_dim_fparam(524-526)get_buffer_dim_fparam(532-534)get_dim_aparam(528-530)get_buffer_dim_aparam(536-538)deepmd/pd/model/task/fitting.py (7)
get_buffer_type_map(478-480)set_case_embd(482-489)deserialize(431-439)get_dim_fparam(441-443)get_buffer_dim_fparam(449-451)get_dim_aparam(445-447)get_buffer_dim_aparam(453-455)deepmd/utils/path.py (1)
DPPath(28-158)deepmd/pd/infer/deep_eval.py (3)
get_dim_fparam(215-219)has_default_fparam(227-233)get_dim_aparam(221-225)
deepmd/pd/infer/deep_eval.py (10)
deepmd/infer/deep_eval.py (18)
output_def(420-421)DeepEval(365-741)DeepEvalBackend(31-362)get_model(351-362)get_model(729-741)get_rcut(150-151)get_rcut(423-425)eval(102-147)get_dim_fparam(162-163)get_dim_fparam(435-437)get_dim_aparam(170-171)get_dim_aparam(443-445)has_default_fparam(165-167)has_default_fparam(439-441)get_var_name(330-332)model_type(302-303)eval_fitting_last_layer(224-264)eval_fitting_last_layer(538-603)deepmd/dpmodel/output_def.py (2)
OutputVariableCategory(130-148)OutputVariableDef(151-239)deepmd/jax/infer/deep_eval.py (7)
DeepEval(60-434)get_model(426-434)eval(189-252)get_dim_fparam(143-145)get_dim_aparam(147-149)model_type(152-166)_get_output_shape(398-420)deepmd/pt/infer/deep_eval.py (11)
DeepEval(85-841)get_model(723-731)get_rcut(201-203)eval(306-380)get_dim_fparam(213-215)get_dim_aparam(217-219)has_default_fparam(221-227)get_var_name(232-239)model_type(242-260)_get_output_shape(623-647)eval_fitting_last_layer(788-841)deepmd/pd/model/network/network.py (1)
TypeEmbedNetConsistent(127-341)deepmd/utils/econf_embd.py (1)
sort_element_type(273-283)deepmd/utils/model_branch_dict.py (1)
get_model_dict(12-42)deepmd/pd/model/atomic_model/dp_atomic_model.py (10)
get_rcut(154-156)get_buffer_rcut(166-168)get_buffer_type_map(162-164)get_buffer_dim_fparam(402-404)get_dim_fparam(398-400)get_dim_aparam(410-412)get_buffer_dim_aparam(414-416)has_default_fparam(406-408)eval_fitting_last_layer(142-144)set_eval_fitting_last_layer_hook(135-140)deepmd/pd/model/model/make_model.py (8)
get_rcut(556-558)get_buffer_rcut(564-566)get_buffer_type_map(568-570)get_buffer_dim_fparam(532-534)get_dim_fparam(524-526)get_dim_aparam(528-530)get_buffer_dim_aparam(536-538)model_output_type(87-98)deepmd/pd/train/wrapper.py (1)
ModelWrapper(22-215)
deepmd/pd/model/atomic_model/base_atomic_model.py (13)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_type_map(162-164)has_default_fparam(406-408)deepmd/pd/model/descriptor/se_a.py (3)
get_buffer_type_map(154-156)get_ntypes(146-148)get_ntypes(554-556)deepmd/pd/model/model/make_model.py (1)
get_buffer_type_map(568-570)deepmd/pd/model/descriptor/dpa1.py (2)
get_buffer_type_map(356-358)get_ntypes(348-350)deepmd/pd/model/descriptor/dpa3.py (2)
get_buffer_type_map(260-262)get_ntypes(252-254)deepmd/pd/model/model/ener_model.py (1)
get_buffer_type_map(37-39)deepmd/pd/model/descriptor/se_t_tebd.py (3)
get_buffer_type_map(216-218)get_ntypes(208-210)get_ntypes(649-651)deepmd/pd/infer/deep_eval.py (2)
has_default_fparam(227-233)get_ntypes(207-209)deepmd/pd/model/descriptor/repflows.py (1)
get_ntypes(383-385)deepmd/pd/model/descriptor/dpa2.py (1)
get_ntypes(359-361)deepmd/pd/model/model/model.py (2)
get_ntypes(53-55)get_buffer_ntypes(57-59)deepmd/pd/model/descriptor/repformers.py (1)
get_ntypes(348-350)deepmd/pd/model/descriptor/se_atten.py (1)
get_ntypes(303-305)
deepmd/pd/utils/stat.py (2)
deepmd/pt/utils/stat.py (3)
make_stat_input(39-94)_compute_model_predict(161-197)rmse(536-537)deepmd/pd/model/atomic_model/base_atomic_model.py (1)
model_forward(513-544)
🪛 Ruff (0.13.1)
deepmd/pd/infer/deep_eval.py
112-112: Unused method argument: no_jit
(ARG002)
113-113: Unused method argument: kwargs
(ARG002)
241-241: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
724-724: Avoid specifying long messages outside the exception class
(TRY003)
deepmd/pd/model/atomic_model/base_atomic_model.py
319-319: Unused method argument: model_with_new_type_stat
(ARG002)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (false)
🔇 Additional comments (21)
deepmd/pd/entrypoints/main.py (7)
10-10: LGTM: Import addition is appropriate.The
Anyimport from typing is necessary for the new type annotations in the function signatures.
83-92: LGTM: Comprehensive type annotations added.The function signature has been properly annotated with detailed parameter types and return type. The annotations align with the actual usage in the codebase.
101-106: LGTM: Inner function annotations are appropriate.The parameter and return type annotations for the inner function are properly specified and consistent with the parent function's usage.
113-117: LGTM: Enhanced data system processing with pattern support.The changes correctly utilize the
process_systemsfunction with the newpatternsparameter, extractingrglob_patternsfrom the dataset configuration and passing them appropriately for both training and validation systems.
351-351: LGTM: New atomic virial parameter addition.The
do_atomic_virialparameter addition to the freeze function signature is appropriate for enabling atomic-virial aware static graphs.
384-384: Fix the boolean parameter in input_spec.Based on the previous review feedback, raw boolean values should not be included in
input_spec. The boolean parameter needs to be handled differently - either removed from input_spec (if it's meant to be frozen as a constant) or converted to a proper InputSpec.Apply this fix to remove the boolean from both input_spec lists:
None, # fparam None, # aparam - do_atomic_virial, # do_atomic_virial + # do_atomic_virial frozen as constant - not in input_specAnd apply the same change to forward_lower's input_spec.
Also applies to: 406-406
419-438: Add model_output_type to the paddle.jit.to_static wrappers (deepmd/pd/entrypoints/main.py)get_buffer_* methods are present across PD models (hasattr guards missing ones). model_output_type is implemented in PD model files and used by deepmd/pd/infer/deep_eval.py for evaluator selection — include it in the wrapped list.
for method_name in [ "get_buffer_rcut", "get_buffer_type_map", "get_buffer_dim_fparam", "get_buffer_dim_aparam", "get_buffer_intensive", "get_buffer_sel_type", "get_buffer_numb_dos", "get_buffer_task_dim", + "model_output_type", ]:deepmd/pd/utils/stat.py (2)
65-72: LGTM: Correct handling for default fparamRemoving fparam/find_fparam when find_fparam == 0.0 matches the PT logic and avoids spurious fields during stats.
228-229: LGTM: Return type clarified for _fill_stat_with_globalAccurate use of Union[np.ndarray, None] matches behavior.
source/tests/pd/model/test_atomic_model_atomic_stat.py (2)
118-121: Type annotations on setUp/tearDown look goodClearer intent; matches surrounding tests.
241-243: Good: assert messages improve failure diagnosticsAdding err_msg helps pinpoint mismatches.
deepmd/pd/model/descriptor/se_a.py (2)
554-557: Dynamic/static ntypes resolution LGTMReturning buffer in static mode is correct.
98-103: Always register buffer_type_map (even when type_map is None)Prevents attribute errors in static mode when accessing get_buffer_type_map().
Apply:
- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)], dtype="int32"), + ) + else: + self.register_buffer( + "buffer_type_map", paddle.to_tensor([], dtype="int32") + )deepmd/pd/infer/deep_eval.py (6)
719-726: eval_typeebd: LGTMCorrectly aggregates TypeEmbedNetConsistent outputs.
2-2: Reminder: run ruff check/format on Python changesPer repo guidelines, run “ruff check .” and “ruff format .”.
You can run locally or in CI to catch ARG002/B009/TRY003 flagged above.
498-499: Fix static box shape: should be [nframes, 9]Freeze uses InputSpec([-1, 9]); feeding [nframes, 3, 3] will mis-shape inputs.
- self.box_handle.copy_from_cpu(cells.reshape([nframes, 3, 3])) + self.box_handle.copy_from_cpu(cells.reshape([nframes, 9]))
586-597: Block spin path for static modelsStatic (.json) spin inference isn’t supported; fail fast.
def _eval_model_spin( @@ - model = self.dp.to(DEVICE) + if self.static_model: + raise NotImplementedError( + "Spin inference for static (.json) models is not supported; use a .pd model." + ) + model = self.dp.to(DEVICE)
194-201: Robust has_spin detection; avoid calling non-callablesPrevents CodeQL “non-callable called”.
- self._has_spin = ( - getattr(self.dp.model["Default"], "has_spin", False) - if isinstance(self.dp, ModelWrapper) - else False - ) - if callable(self._has_spin): - self._has_spin = False + has_spin_attr = ( + getattr(self.dp.model["Default"], "has_spin", None) + if isinstance(self.dp, ModelWrapper) + else None + ) + if callable(has_spin_attr): + self._has_spin = bool(has_spin_attr()) + else: + self._has_spin = bool(has_spin_attr) if has_spin_attr is not None else False
544-558: Static do_atomic_virial control: add optional input handle or bake constantCurrently computed conditionally but never fed into static graph.
Option A (input handle if exported):
+ # Optional control for atomic virial + if "do_atomic_virial" in self.predictor.get_input_names(): + self.do_atomic_virial_handle = self.predictor.get_input_handle( + "do_atomic_virial" + ) @@ - self.predictor.run() + if hasattr(self, "do_atomic_virial_handle"): + self.do_atomic_virial_handle.copy_from_cpu( + np.asarray([do_atomic_virial], dtype="bool") + ) + self.predictor.run()Option B: export a fixed setting and document it.
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
87-95: Handle type_map=None by registering an empty bufferPrevents missing-attribute errors in static graphs.
- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)], dtype="int32"), + ) + else: + self.register_buffer("buffer_type_map", paddle.empty([0], dtype="int32"))deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
170-173: Guard get_buffer_sel for descriptors without this methodSome descriptors don’t expose get_buffer_sel(); provide a safe fallback returning a tensor from self.sel.
- def get_buffer_sel(self) -> paddle.Tensor: - """Get the neighbor selection.""" - return self.descriptor.get_buffer_sel() + def get_buffer_sel(self) -> paddle.Tensor: + """Get the neighbor selection.""" + if hasattr(self.descriptor, "get_buffer_sel"): + return self.descriptor.get_buffer_sel() + # Fallback: build a tensor from self.sel + return paddle.to_tensor(self.sel, dtype="int64")
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (7)
deepmd/pd/model/atomic_model/base_atomic_model.py (2)
297-302: Remove stray debug comment.The
# raiseline looks like a leftover and may confuse readers.Apply this diff:
- ret_dict["mask"] = atom_mask - # raise + ret_dict["mask"] = atom_mask
610-615: Support static/JIT mode: use buffer in non-dynamic mode.Align with descriptors and other modules: return a buffer-backed value when not in dynamic mode to improve static export compatibility.
Apply this diff:
- def get_ntypes(self): - return len(self.type_map) + def get_ntypes(self): + return len(self.type_map) if paddle.in_dynamic_mode() else self.buffer_ntypesdeepmd/pd/model/descriptor/dpa2.py (1)
786-789: Prefer accessor methods for consistency.Use
get_rcut()andget_nsel()to build the key, mirroring the call below forrepformers.Apply this diff:
- nlist_dict[ - get_multiple_nlist_key(self.repinit.rcut, sum(self.repinit.sel)) - ], + nlist_dict[ + get_multiple_nlist_key(self.repinit.get_rcut(), self.repinit.get_nsel()) + ],deepmd/pd/model/descriptor/dpa3.py (2)
260-271: Docstring ok; ensure buffer exists and dtype consistency.This accessor is fine; it depends on buffer_type_map being registered even when type_map is None (see earlier comment). Consider standardizing int32 dtype across all modules for this buffer.
215-217: Unify integer dtype for buffer_ntypes (optional).You use int64 here; elsewhere (e.g., DPAtomicModel) uses int32. Aligning dtypes reduces friction in cross-module static consumers.
deepmd/pd/model/descriptor/se_a.py (2)
466-469: Specify dtype for buffer_rcut/buffer_rcut_smth to match precision.Without dtype, Paddle defaults may be float32. Use the class precision for consistency.
Apply this diff:
- self.register_buffer("buffer_rcut", paddle.to_tensor(self.rcut)) + self.register_buffer("buffer_rcut", paddle.to_tensor(self.rcut, dtype=self.prec)) ... - self.register_buffer("buffer_rcut_smth", paddle.to_tensor(self.rcut_smth)) + self.register_buffer("buffer_rcut_smth", paddle.to_tensor(self.rcut_smth, dtype=self.prec))
479-481: Unify integer dtype for buffer_ntypes (optional).Consider using a consistent dtype (int32 or int64) across modules for buffer_ntypes.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
deepmd/pd/model/atomic_model/base_atomic_model.py(13 hunks)deepmd/pd/model/atomic_model/dp_atomic_model.py(10 hunks)deepmd/pd/model/descriptor/dpa1.py(3 hunks)deepmd/pd/model/descriptor/dpa2.py(5 hunks)deepmd/pd/model/descriptor/dpa3.py(4 hunks)deepmd/pd/model/descriptor/repflows.py(3 hunks)deepmd/pd/model/descriptor/repformers.py(3 hunks)deepmd/pd/model/descriptor/se_a.py(7 hunks)deepmd/pd/model/descriptor/se_atten.py(4 hunks)deepmd/pd/model/descriptor/se_t_tebd.py(6 hunks)deepmd/pd/model/model/ener_model.py(1 hunks)deepmd/pd/model/model/make_model.py(2 hunks)deepmd/pd/model/model/model.py(1 hunks)deepmd/pd/model/task/fitting.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (5)
- deepmd/pd/model/descriptor/dpa1.py
- deepmd/pd/model/model/model.py
- deepmd/pd/model/task/fitting.py
- deepmd/pd/model/descriptor/repformers.py
- deepmd/pd/model/descriptor/se_atten.py
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/model/descriptor/repflows.pydeepmd/pd/model/descriptor/dpa2.pydeepmd/pd/model/descriptor/se_t_tebd.pydeepmd/pd/model/descriptor/se_a.pydeepmd/pd/model/model/make_model.pydeepmd/pd/model/atomic_model/base_atomic_model.pydeepmd/pd/model/descriptor/dpa3.pydeepmd/pd/model/model/ener_model.pydeepmd/pd/model/atomic_model/dp_atomic_model.py
🧬 Code graph analysis (9)
deepmd/pd/model/descriptor/repflows.py (4)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_rcut(174-176)get_buffer_sel(178-180)deepmd/pd/model/descriptor/dpa3.py (3)
get_buffer_rcut(232-234)get_buffer_sel(248-250)get_ntypes(252-254)deepmd/pd/model/descriptor/repformers.py (2)
get_buffer_rcut(332-334)get_ntypes(348-350)deepmd/pd/model/descriptor/se_atten.py (2)
get_buffer_rcut(287-289)get_ntypes(303-305)
deepmd/pd/model/descriptor/dpa2.py (3)
deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
get_buffer_rcut(174-176)deepmd/pd/model/descriptor/repflows.py (1)
get_buffer_rcut(363-365)deepmd/pd/utils/nlist.py (1)
get_multiple_nlist_key(332-336)
deepmd/pd/model/descriptor/se_t_tebd.py (5)
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-147)deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_type_map(162-172)get_buffer_rcut(174-176)deepmd/pd/model/descriptor/se_a.py (5)
get_buffer_type_map(154-164)get_buffer_rcut(130-132)get_buffer_rcut(546-548)get_buffer_rcut_smth(134-136)get_buffer_rcut_smth(550-552)deepmd/pd/model/model/make_model.py (2)
get_buffer_type_map(568-578)get_buffer_rcut(564-566)deepmd/pd/model/descriptor/repflows.py (1)
get_buffer_rcut(363-365)
deepmd/pd/model/descriptor/se_a.py (5)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_rcut(174-176)get_buffer_type_map(162-172)deepmd/pd/model/descriptor/repflows.py (1)
get_buffer_rcut(363-365)deepmd/pd/model/descriptor/se_atten.py (2)
get_buffer_rcut(287-289)get_buffer_rcut_smth(291-293)deepmd/pd/model/model/make_model.py (2)
get_buffer_rcut(564-566)get_buffer_type_map(568-578)deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-147)
deepmd/pd/model/model/make_model.py (3)
deepmd/pd/model/atomic_model/dp_atomic_model.py (4)
get_buffer_dim_fparam(410-412)get_buffer_dim_aparam(422-424)get_buffer_rcut(174-176)get_buffer_type_map(162-172)deepmd/pd/model/task/fitting.py (3)
get_buffer_dim_fparam(449-451)get_buffer_dim_aparam(453-455)get_buffer_type_map(478-488)deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-147)
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
deepmd/pd/model/atomic_model/dp_atomic_model.py (2)
get_buffer_type_map(162-172)has_default_fparam(414-416)
deepmd/pd/model/descriptor/dpa3.py (4)
deepmd/pd/model/atomic_model/dp_atomic_model.py (4)
get_buffer_rcut(174-176)get_buffer_sel(178-180)get_sel(158-160)get_buffer_type_map(162-172)deepmd/pd/model/descriptor/repflows.py (4)
get_buffer_rcut(363-365)get_buffer_sel(379-381)get_sel(375-377)get_ntypes(383-385)deepmd/pd/model/atomic_model/base_atomic_model.py (3)
get_ntypes(610-611)get_type_map(133-135)get_buffer_type_map(137-147)source/api_cc/src/DeepPotPD.cc (2)
get_type_map(658-667)get_type_map(658-658)
deepmd/pd/model/model/ener_model.py (2)
deepmd/pd/model/atomic_model/base_atomic_model.py (1)
get_buffer_type_map(137-147)deepmd/pd/model/model/make_model.py (1)
get_buffer_type_map(568-578)
deepmd/pd/model/atomic_model/dp_atomic_model.py (9)
deepmd/pd/model/atomic_model/base_atomic_model.py (3)
get_buffer_type_map(137-147)deserialize(371-388)has_default_fparam(157-159)deepmd/pd/model/descriptor/dpa1.py (3)
get_buffer_type_map(356-366)get_buffer_rcut(328-330)deserialize(554-590)deepmd/pd/model/descriptor/dpa3.py (4)
get_buffer_type_map(260-270)get_buffer_rcut(232-234)get_buffer_sel(248-250)deserialize(454-484)deepmd/pd/model/descriptor/se_a.py (4)
get_buffer_type_map(154-164)get_buffer_rcut(130-132)get_buffer_rcut(546-548)deserialize(382-398)deepmd/pd/model/descriptor/se_t_tebd.py (3)
get_buffer_type_map(216-226)get_buffer_rcut(641-643)deserialize(398-427)deepmd/pd/model/model/make_model.py (8)
get_buffer_type_map(568-578)get_buffer_rcut(564-566)set_case_embd(521-522)deserialize(518-519)get_dim_fparam(524-526)get_buffer_dim_fparam(532-534)get_dim_aparam(528-530)get_buffer_dim_aparam(536-538)deepmd/pd/model/task/fitting.py (7)
get_buffer_type_map(478-488)set_case_embd(490-497)deserialize(431-439)get_dim_fparam(441-443)get_buffer_dim_fparam(449-451)get_dim_aparam(445-447)get_buffer_dim_aparam(453-455)deepmd/dpmodel/atomic_model/dp_atomic_model.py (4)
set_case_embd(70-75)get_dim_fparam(230-232)has_default_fparam(238-240)get_dim_aparam(234-236)deepmd/pd/infer/deep_eval.py (3)
get_dim_fparam(215-219)has_default_fparam(227-233)get_dim_aparam(221-225)
🪛 Ruff (0.13.1)
deepmd/pd/model/atomic_model/base_atomic_model.py
327-327: Unused method argument: model_with_new_type_stat
(ARG002)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (28)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
🔇 Additional comments (27)
deepmd/pd/model/atomic_model/base_atomic_model.py (2)
137-147: LGTM: public accessor for buffer-backed type_map.The docstring and return type match usage across the codebase.
87-95: Always registerbuffer_type_map; handletype_map=None.Currently the buffer is only registered when
type_mapis not None. Callers ofget_buffer_type_map()(Line 147) will hit AttributeError if the buffer is absent, and static paths expect buffers to always exist. Also makeself.ntypesrobust whentype_mapis None.Apply this diff:
- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) - self.ntypes = len(self.type_map) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)], dtype="int32"), + ) + else: + # ensure buffer exists for static/JIT even when type_map is absent + self.register_buffer("buffer_type_map", paddle.to_tensor([], dtype="int32")) + self.ntypes = len(type_map) if type_map is not None else 0deepmd/pd/model/model/ener_model.py (1)
37-47: LGTM: passthroughget_buffer_type_map()matches the new buffer-access pattern.Exposes the buffer-backed type map at the model level without altering behavior.
deepmd/pd/model/descriptor/dpa2.py (3)
326-328: LGTM: persistbuffer_ntypesas int64 for static inference.Matches usage elsewhere.
343-350: LGTM: delegate buffer-backed rcut/rcut_smth to repinit.
268-273: Always registerbuffer_type_map; handletype_map=None.Without the else branch, static paths querying
get_buffer_type_map()will raise AttributeError.Apply this diff:
- if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)], dtype="int32"), + ) + else: + self.register_buffer("buffer_type_map", paddle.to_tensor([], dtype="int32"))deepmd/pd/model/descriptor/repflows.py (3)
241-254: Fix buffer dtypes and add missingbuffer_rcut_smth.
buffer_selshould be integer-typed.buffer_rcut_smthis missing; callers expectget_buffer_rcut_smth().
[ suggest_recommended_refactor ]Apply this diff:
self.sel = sel - self.register_buffer("buffer_sel", paddle.to_tensor(sel)) + self.register_buffer("buffer_sel", paddle.to_tensor(sel, dtype="int64")) self.rcut = e_rcut self.register_buffer("buffer_rcut", paddle.to_tensor(self.e_rcut)) + # keep symmetry with rcut for static inference users + self.register_buffer("buffer_rcut_smth", paddle.to_tensor(self.e_rcut_smth))
363-366: Exposeget_buffer_rcut_smth()to match callers.Prevents AttributeError in modules that need the smoothed cutoff radius buffer.
Apply this diff:
def get_buffer_rcut(self) -> paddle.Tensor: """Returns the cut-off radius as a buffer-style Tensor.""" return self.buffer_rcut + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where smoothing starts as a buffer-style Tensor.""" + return self.buffer_rcut_smth
379-386: LGTM: dynamic/staticget_ntypespattern.deepmd/pd/model/descriptor/se_t_tebd.py (4)
212-215: Reconstructtype_mapfrom buffer in non-dynamic mode.Static export won’t carry Python lists; provide a robust fallback.
Apply this diff:
def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" - return self.type_map + if paddle.in_dynamic_mode(): + return self.type_map + if not hasattr(self, "buffer_type_map") or self.buffer_type_map.numel() == 0: + return [] + # decode ASCII codes into a space-separated string and split + codes = [int(x) for x in self.buffer_type_map.numpy().tolist()] + return "".join(chr(c) for c in codes).split(" ")
641-648: LGTM: buffer-backed rcut/rcut_smth accessors in the block.
657-660: LGTM: dynamic/staticget_ntypespattern.
168-173: Always registerbuffer_type_map; handletype_map=None.Ensure the buffer exists for static/JIT paths even when
type_mapis not provided.Apply this diff:
self.type_map = type_map - if type_map is not None: - self.register_buffer( - "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), - ) + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)], dtype="int32"), + ) + else: + self.register_buffer("buffer_type_map", paddle.to_tensor([], dtype="int32"))deepmd/pd/model/model/make_model.py (2)
532-539: Expose buffer dims: verify downstream expects Tensors (not ints).APIs look good. DeepEval.get_dim_* is annotated to return int but may receive a paddle.Tensor in static mode. Ensure callers cast to Python ints if needed (e.g., int(t.item()) or int(t.numpy()[0])).
564-579: LGTM: buffer accessors for rcut and type_map.Consistent with atomic model APIs; helps static/JIT paths.
deepmd/pd/model/descriptor/dpa3.py (4)
231-239: LGTM: buffer rcut/smth accessors.Delegation to repflows’ buffers is correct.
252-255: LGTM: dynamic vs static ntypes.Appropriate use of in_dynamic_mode to choose between Python value and buffer.
256-259: Handle None and static-mode decoding in get_type_map.Ensure safe behavior when type_map is None and when running in static mode where buffer decoding is required.
Apply this diff:
def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" - return self.type_map + if paddle.in_dynamic_mode(): + return self.type_map if self.type_map is not None else [] + # static mode: decode from buffer for JIT/predictor compatibility + if getattr(self, "buffer_type_map", None) is None or self.buffer_type_map.numel() == 0: + return [] + return "".join([chr(int(x)) for x in self.buffer_type_map.numpy()]).split(" ")
248-251: Return Tensor from get_buffer_sel (not Python list).This currently returns a Python list and will break static export/predictor paths that expect a Tensor; DPAtomicModel.get_buffer_sel relies on this being a Tensor.
Apply this diff:
def get_buffer_sel(self) -> paddle.Tensor: """Returns the number of selected atoms for each type as a buffer-style Tensor.""" - return self.repflows.get_sel() + return self.repflows.get_buffer_sel()deepmd/pd/model/descriptor/se_a.py (3)
130-137: LGTM: buffer rcut/smth accessors.Matches block-level buffers.
154-165: LGTM, contingent on buffer registration fix.Accessor is fine once buffer_type_map is always registered.
98-103: Always register buffer_type_map; handle None and set dtype.Static paths require the buffer attribute to exist. Register an empty int tensor when type_map is None.
Apply this diff:
self.type_map = type_map - if type_map is not None: + if type_map is not None: self.register_buffer( "buffer_type_map", - paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + paddle.to_tensor([ord(c) for c in " ".join(type_map)], dtype="int32"), ) + else: + self.register_buffer("buffer_type_map", paddle.to_tensor([], dtype="int32"))deepmd/pd/model/atomic_model/dp_atomic_model.py (5)
73-84: Good: replaced 3.10|with typing.Union.Fixes Python version compatibility.
162-173: LGTM: buffer_type_map accessor.Matches buffer encoding used across the codebase.
245-256: Deserialize via BaseAtomicModel: verify this still constructs DPAtomicModel.Assuming BaseAtomicModel.deserialize calls cls(**data), this is fine. Please confirm no regression in loading models trained pre-PR.
358-405: LGTM: sampled_func typing and optional output-stat computation.The flag enables input-only stats; aligns with intended use.
410-417: LGTM: buffer dim getters + has_default_fparam.Matches fitting interface and static consumers.
| @coderabbitai resolve |
✅ Actions performedComments resolved. |
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu> Signed-off-by: HydrogenSulfate <490868991@qq.com>
There was a problem hiding this comment.
Actionable comments posted: 2
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pd/infer/deep_eval.py(20 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Always run
ruff check .andruff format .before committing changes to Python code
Files:
deepmd/pd/infer/deep_eval.py
🧠 Learnings (1)
📓 Common learnings
Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-09-19T04:25:12.408Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. Learnt from: njzjz PR: deepmodeling/deepmd-kit#4144 File: source/api_cc/tests/test_deeppot_dpa_pt.cc:166-246 Timestamp: 2024-10-08T15:32:11.479Z Learning: Refactoring between test classes `TestInferDeepPotDpaPt` and `TestInferDeepPotDpaPtNopbc` is addressed in PR #3905. 🧬 Code graph analysis (1)
deepmd/pd/infer/deep_eval.py (5)
deepmd/pd/model/network/network.py (1)
TypeEmbedNetConsistent(127-341)deepmd/utils/econf_embd.py (1)
sort_element_type(273-283)deepmd/utils/model_branch_dict.py (1)
get_model_dict(12-42)deepmd/pd/model/atomic_model/dp_atomic_model.py (7)
get_buffer_type_map(162-172)get_buffer_dim_fparam(410-412)get_dim_fparam(406-408)get_dim_aparam(418-420)get_buffer_dim_aparam(422-424)has_default_fparam(414-416)eval_fitting_last_layer(142-144)deepmd/pd/train/wrapper.py (1)
ModelWrapper(22-215)
🪛 Ruff (0.13.1)
deepmd/pd/infer/deep_eval.py
112-112: Unused method argument: no_jit
(ARG002)
113-113: Unused method argument: kwargs
(ARG002)
241-241: Do not call getattr with a constant attribute value. It is not any safer than normal property access.
Replace getattr with attribute access
(B009)
724-724: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (29)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Analyze (python)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.9)
🔇 Additional comments (3)
deepmd/pd/infer/deep_eval.py (3)
157-165: Don’t hard-code GPU usage for predictorsCalling
config.enable_use_gpu(4096, 0)unconditionally forces every static inference run onto GPU 0 with exactly 4 GB reserved. On CPU-only hosts (or when GPU 0 is unavailable) predictor creation now fails outright, making the newly added.jsonsupport unusable in common environments. Please make GPU usage opt-in/parametrized and retain a CPU fallback.- config.enable_use_gpu(4096, 0) + if kwargs.get("use_gpu", False): + gpu_mem = kwargs.get("gpu_mem", 4096) + gpu_id = kwargs.get("gpu_id", 0) + config.enable_use_gpu(gpu_mem, gpu_id)
194-201: Restore callablehas_spindetectionThe new
if callable(self._has_spin): self._has_spin = Falsebranch now forces_has_spintoFalsewhenever the model exposes the usualhas_spin()method. Spin-capable models are therefore reported as spinless and upstream code will skip supplying spin tensors, silently breaking results. Call the attribute when it is callable instead of discarding it.- self._has_spin = ( - getattr(self.dp.model["Default"], "has_spin", False) - if isinstance(self.dp, ModelWrapper) - else False - ) - if callable(self._has_spin): - self._has_spin = False + has_spin_attr = ( + getattr(self.dp.model["Default"], "has_spin", None) + if isinstance(self.dp, ModelWrapper) + else None + ) + if callable(has_spin_attr): + self._has_spin = bool(has_spin_attr()) + else: + self._has_spin = bool(has_spin_attr) if has_spin_attr is not None else False
654-658: Use Paddle-compatible tensor→NumPy conversion in spin path
paddle.Tensordoes not expose.cpu(). The new.detach().cpu().numpy()chain will therefore raiseAttributeErrorand break every spin evaluation. Reuseto_numpy_array()(already imported) or call.numpy()directly to keep this path working.- out = batch_output[pd_name].reshape(shape).detach().cpu().numpy() + out = to_numpy_array(batch_output[pd_name].reshape(shape))
2 participants
Summary:
DP(model=paddle_inference.json)for ase, and add an unitest.DeepPotPD.ccThis pull request introduces significant enhancements and refactoring to the Paddle backend implementation of
DeepEvalindeepmd/pd/infer/deep_eval.py, along with minor improvements to model freezing indeepmd/pd/entrypoints/main.py. The most important changes are the addition of support for static models loaded from.jsonfiles, expanded model type handling, and improved output shape and evaluation logic for various model branches and output variable categories.Static model and inference support
.jsonfiles, including Paddle inference engine integration and input/output handle management for efficient prediction. (deepmd/pd/infer/deep_eval.py) [1] [2] [3] [4] [5]Expanded model type and branch handling
deepmd/pd/infer/deep_eval.py) [1] [2]Output variable and evaluation improvements
DERV_R_DERV_R), and refactored evaluation logic to handle both static and dynamic models, including proper output conversion. (deepmd/pd/infer/deep_eval.py) [1] [2] [3]_eval_model_spinmethod to support spin-dependent model evaluation and output extraction. (deepmd/pd/infer/deep_eval.py)Type embedding evaluation
eval_typeebdto extract and concatenate type embedding network outputs from the loaded model. (deepmd/pd/infer/deep_eval.py)Model freezing improvements
freezefunction to support freezing additional model methods and atomic virial computation, and changed input specifications for better compatibility. (deepmd/pd/entrypoints/main.py) [1] [2] [3] [4]Let me know if you'd like a walkthrough of any specific new functionality or code sections!
Summary by CodeRabbit
New Features
Improvements
Documentation
Tests