fix(pd): update model interface for dpa3 paddle backend#5319
Open
HydrogenSulfate wants to merge 4 commits intodeepmodeling:masterfrom
Open
fix(pd): update model interface for dpa3 paddle backend#5319HydrogenSulfate wants to merge 4 commits intodeepmodeling:masterfrom
HydrogenSulfate wants to merge 4 commits intodeepmodeling:masterfrom
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Updates the Paddle (pd) backend DPA3 model/descriptor interfaces to support optional charge/spin embeddings and aligns parts of the Paddle atomic model behavior with other backends, with accompanying test updates.
Changes:
- Add
add_chg_spin_ebd+fparamsupport to the Paddle DPA3 descriptor (including (de)serialization). - Extend Paddle DPA3 tests to cover the charge/spin embedding path.
- Update Paddle atomic model docs and adjust statistics computation flow (including observed-type handling hook).
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 7 comments.
| File | Description |
|---|---|
source/tests/pd/model/test_dpa3.py |
Expands consistency testing to include add_chg_spin_ebd and passes fparam through both Paddle and dpmodel code paths. |
deepmd/pd/model/descriptor/dpa3.py |
Adds charge/spin embedding components, serializes them, and threads fparam into forward(). |
deepmd/pd/model/atomic_model/dp_atomic_model.py |
Enhances class docstring, captures add_chg_spin_ebd, modifies stat computation flow, and adds an observed-type collection call. |
Comments suppressed due to low confidence (1)
deepmd/pd/model/descriptor/dpa3.py:566
forward()gained anfparamargument and now conditionally requires it whenadd_chg_spin_ebdis enabled, but the method docstring doesn’t mentionfparam(shape/meaning) or the requirement. Please update the docstring (and ideally validate with a clear error) so failures are actionable without inspecting the code.
def forward(
self,
extended_coord: paddle.Tensor,
extended_atype: paddle.Tensor,
nlist: paddle.Tensor,
mapping: paddle.Tensor | None = None,
comm_dict: list[paddle.Tensor] | None = None,
fparam: paddle.Tensor | None = None,
) -> tuple[
paddle.Tensor,
paddle.Tensor | None,
paddle.Tensor | None,
paddle.Tensor | None,
paddle.Tensor | None,
]:
"""Compute the descriptor.
Parameters
----------
extended_coord
The extended coordinates of atoms. shape: nf x (nallx3)
extended_atype
The extended aotm types. shape: nf x nall
nlist
The neighbor list. shape: nf x nloc x nnei
mapping
The index mapping, mapps extended region index to local region.
comm_dict
The data needed for communication for parallel inference.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+82
to
+84
| self.add_chg_spin_ebd: bool = getattr( | ||
| self.descriptor, "add_chg_spin_ebd", False | ||
| ) |
| self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) | ||
| self.compute_fitting_input_stat(wrapped_sampler, stat_file_path) | ||
| self.fitting_net.compute_input_stats( | ||
| wrapped_sampler, stat_file_path=stat_file_path |
Comment on lines
+426
to
+428
| self._collect_and_set_observed_type( | ||
| wrapped_sampler, stat_file_path, preset_observed_type | ||
| ) |
| def deserialize(cls, data: dict) -> "DPAtomicModel": | ||
| data = data.copy() | ||
| check_version_compatibility(data.pop("@version", 1), 2, 1) | ||
| check_version_compatibility(data.pop("@version", 1), 2, 2) |
| self.fitting_net = fitting | ||
| if hasattr(self.fitting_net, "reinit_exclude"): | ||
| self.fitting_net.reinit_exclude(self.atom_exclude_types) | ||
| self.type_map = type_map |
| only when needed. Since the sampling process can be slow and memory-intensive, | ||
| the lazy function helps by only sampling once. | ||
| Each element, ``merged[i]``, is a data dictionary containing | ||
| ``keys``: ``np.ndarray`` originating from the ``i``-th data system. |
Comment on lines
120
to
125
| use_econf_tebd: bool = False, | ||
| use_tebd_bias: bool = False, | ||
| use_loc_mapping: bool = True, | ||
| type_map: list[str] | None = None, | ||
| add_chg_spin_ebd: bool = False, | ||
| ) -> None: |
Contributor
📝 WalkthroughWalkthroughAdds optional charge/spin system embeddings to DPA3 and threads fparam/default-fparam and observed-type handling across descriptors, fitting, atomic-model, CM, and stat I/O; expands descriptor/descriptor-block forward signatures to accept/propagate fparam and returns extended optional outputs. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant CM
participant AtomicModel
participant Descriptor
participant Fitting
participant StatIO
rect rgba(200,200,255,0.5)
User->>CM: forward_common(coord, atype, nlist, ..., coord_corr_for_virial?)
CM->>AtomicModel: forward_common_lower(..., extended_coord_corr?)
AtomicModel->>Descriptor: forward(coord, atype, nlist, fparam?)
Descriptor-->>AtomicModel: node_embed (+ optional chg/spin system embedding)
AtomicModel->>Fitting: eval_fitting_input(node_embed, fparam?)
Fitting-->>AtomicModel: outputs (energy, forces, virial)
AtomicModel-->>CM: aggregated model outputs
CM->>User: return outputs (with virial correction)
end
rect rgba(200,255,200,0.5)
User->>CM: compute_or_load_stat(sampled_func, stat_file, preset_observed_type?)
CM->>AtomicModel: compute_or_load_stat(..., preset_observed_type)
AtomicModel->>AtomicModel: _make_wrapped_sampler(sampled_func)
AtomicModel->>Descriptor: sample calls (wrapped) -> returns samples (may include fparam)
AtomicModel->>StatIO: save/load stat files (fparam/aparam/out_bias)
AtomicModel->>AtomicModel: _collect_and_set_observed_type(preset_or_computed)
StatIO-->>AtomicModel: persisted buffers restored or written
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~40 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
📝 Coding Plan
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 Tip CodeRabbit can use OpenGrep to find security vulnerabilities and bugs across 17+ programming languages.OpenGrep is compatible with Semgrep configurations. Add an |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pd/model/atomic_model/dp_atomic_model.py`:
- Around line 420-422: The code currently calls
self.fitting_net.compute_input_stats(...) directly which bypasses the wrapper
compute_fitting_input_stat() (and ignores the model's data_stat_protect flag and
any subclass overrides); change the call to invoke
self.compute_fitting_input_stat(wrapped_sampler, stat_file_path=stat_file_path,
protection=self.data_stat_protect) (or otherwise call compute_fitting_input_stat
with the same arguments) so the protection parameter and any overridden behavior
are respected; look for references to self.fitting_net.compute_input_stats,
compute_fitting_input_stat, and data_stat_protect in dp_atomic_model.py to
update the call site.
- Around line 79-84: DPAtomicModel currently records add_chg_spin_ebd but never
wires frame parameters through the descriptor or exposes descriptor fparam
dimensions, causing mismatched input shape and an assert in DescrptDPA3.forward;
update DPAtomicModel.forward_atomic to pass the fparam argument into
self.descriptor(...) when calling the descriptor, and change any places that
compute fparam dim (the "*_dim_fparam" interface/attributes derived from
self.fitting_net) to include the descriptor's required fparam size (query
self.descriptor for its fparam dim when add_chg_spin_ebd is True) so the model
advertises the correct input shape and avoids the assert; touch
methods/attributes around self.fitting_net, self.descriptor, forward_atomic, and
the dim-reporting properties to wire and propagate fparam end-to-end.
In `@deepmd/pd/model/descriptor/dpa3.py`:
- Around line 595-606: Replace the bare asserts with explicit validation of
fparam before indexing: in the block gated by add_chg_spin_ebd check that fparam
is not None and has at least two columns (e.g. fparam.ndim==2 and
fparam.shape[1] >= 2), otherwise raise a clear ValueError mentioning the
missing/insufficient frame parameters for charge/spin; only after these checks
perform the conversions to int64 and the calls to chg_embedding, spin_embedding,
mix_cs_mlp, act and the subsequent add to node_ebd_ext so callers receive a
helpful exception instead of AssertionError/IndexError.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 6834e656-eaa6-4f38-aa32-d83117538f24
📒 Files selected for processing (3)
deepmd/pd/model/atomic_model/dp_atomic_model.pydeepmd/pd/model/descriptor/dpa3.pysource/tests/pd/model/test_dpa3.py
Comment on lines
+595
to
+606
| if self.add_chg_spin_ebd: | ||
| assert fparam is not None | ||
| assert self.chg_embedding is not None | ||
| assert self.spin_embedding is not None | ||
| charge = fparam[:, 0].to(dtype=paddle.int64) + 100 | ||
| spin = fparam[:, 1].to(dtype=paddle.int64) | ||
| chg_ebd = self.chg_embedding(charge) | ||
| spin_ebd = self.spin_embedding(spin) | ||
| sys_cs_embd = self.act( | ||
| self.mix_cs_mlp(paddle.cat((chg_ebd, spin_ebd), dim=-1)) | ||
| ) | ||
| node_ebd_ext = node_ebd_ext + sys_cs_embd.unsqueeze(1) |
Contributor
There was a problem hiding this comment.
Validate fparam explicitly before indexing it.
Line 596 uses assert for a public input precondition and then immediately reads fparam[:, 0] / fparam[:, 1]. If this path is called without fparam or with fewer than two frame parameters, callers get a bare AssertionError or IndexError instead of a proper argument error.
🛡️ Proposed fix
if self.add_chg_spin_ebd:
- assert fparam is not None
+ if fparam is None or fparam.shape[-1] < 2:
+ raise ValueError(
+ "`fparam` must have shape [nf, >=2] when add_chg_spin_ebd is enabled."
+ )
assert self.chg_embedding is not None
assert self.spin_embedding is not None
charge = fparam[:, 0].to(dtype=paddle.int64) + 100
spin = fparam[:, 1].to(dtype=paddle.int64)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/descriptor/dpa3.py` around lines 595 - 606, Replace the bare
asserts with explicit validation of fparam before indexing: in the block gated
by add_chg_spin_ebd check that fparam is not None and has at least two columns
(e.g. fparam.ndim==2 and fparam.shape[1] >= 2), otherwise raise a clear
ValueError mentioning the missing/insufficient frame parameters for charge/spin;
only after these checks perform the conversions to int64 and the calls to
chg_embedding, spin_embedding, mix_cs_mlp, act and the subsequent add to
node_ebd_ext so callers receive a helpful exception instead of
AssertionError/IndexError.
Contributor
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 (1)
deepmd/pd/model/descriptor/dpa2.py (1)
937-939:⚠️ Potential issue | 🟠 MajorMissing compression flag check before raising error.
Similar to
DescrptDPA1, theenable_compression()method unconditionally raises "Compression is already enabled" without checkingself.compress(line 334). The error message is misleading when compression is not actually enabled.🐛 Proposed fix
def enable_compression( self, min_nbor_dist: float, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, ) -> None: # do some checks before the mocel compression process + if self.compress: + raise ValueError("Compression is already enabled.") - raise ValueError("Compression is already enabled.") + raise NotImplementedError("Model compression is not yet supported.")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/model/descriptor/dpa2.py` around lines 937 - 939, The enable_compression method currently always raises ValueError("Compression is already enabled.") without checking the actual compression flag; update the method (enable_compression in class DescrptDPA2) to first inspect self.compress and only raise that ValueError if self.compress is already True, otherwise proceed to enable compression (mirror the conditional behavior in DescrptDPA1) so the error message accurately reflects state.
🧹 Nitpick comments (2)
deepmd/pd/model/descriptor/dpa2.py (1)
852-860: LGTM with minor style note.The return statement correctly handles
Nonevalues. There's a minor inconsistency: line 776 uses.astype(dtype=...)while the return uses.to(dtype=...). Both are valid in PaddlePaddle, but you may want to standardize on one style for consistency.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/model/descriptor/dpa2.py` around lines 852 - 860, Standardize Paddle tensor dtype casting in this module: replace the mixed usages so all casts use the same method (either .astype(dtype=env.GLOBAL_PD_FLOAT_PRECISION) or .to(dtype=env.GLOBAL_PD_FLOAT_PRECISION)) for consistency; update the return tuple (g1, rot_mat, g2, h2, sw) to use the chosen form to match the earlier .astype usage on the other tensors in this file (keep the None checks unchanged).deepmd/pd/model/atomic_model/dp_atomic_model.py (1)
374-374: Consider prefixing unused variableswwith underscore.Static analysis indicates
swis unpacked but never used.🧹 Suggested fix
- descriptor, rot_mat, g2, h2, sw = self.descriptor( + descriptor, rot_mat, g2, h2, _sw = self.descriptor(🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/model/atomic_model/dp_atomic_model.py` at line 374, The tuple returned by self.descriptor is unpacked into descriptor, rot_mat, g2, h2, sw but sw is never used; rename the variable to _sw (or prefix it with an underscore) in the unpacking site inside the method using self.descriptor so static analyzers treat it as intentionally unused and avoid warnings.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pd/model/descriptor/dpa1.py`:
- Around line 301-307: The enable_compression method currently raises
"Compression is already enabled" unconditionally; modify enable_compression to
first check the instance flags self.tebd_compress and self.geo_compress (and any
other compression flags used in this class) and only raise that error if both
relevant flags are already True; otherwise set the appropriate flags and proceed
to enable compression. Also apply the same conditional-check pattern to the
other symmetric location that raises the same error (the second occurrence
referring to the same enable_compression logic) so both places validate
self.tebd_compress/self.geo_compress before raising.
In `@deepmd/pd/model/model/make_model.py`:
- Around line 466-470: In the paddle tensor creation call that currently uses
device=nlist.device (the paddle.ones([...], dtype=nlist.dtype,
device=nlist.device) expression involving n_nf, n_nloc, nnei, n_nnei and nlist),
replace the incorrect device keyword with paddle's place keyword (e.g.,
place=nlist.place) so the call becomes paddle.ones(..., dtype=nlist.dtype,
place=nlist.place); also audit nearby paddle.zeros/paddle.ones calls in the same
scope and switch any device=... usages to place=... while keeping dtype and
shape unchanged.
In `@deepmd/pd/model/model/transform_output.py`:
- Around line 258-260: The paddle.zeros call creating virial uses the wrong
keyword argument 'device'; change it to use 'place' (matching Paddle API and the
earlier fix at line ~240) so the call uses place=vv.place (or equivalent)
instead of device=vv.device; update the virial creation in transform_output (the
paddle.zeros(...) that references vldims, derv_c_ext_dims and vv) to use
place=vv.place.
- Around line 239-241: The paddle.zeros call is passing vv.device which doesn't
exist on paddle.Tensor; change the device argument to use vv.place instead.
Locate the paddle.zeros invocation that builds force (the line creating force =
paddle.zeros(vldims + derv_r_ext_dims, dtype=vv.dtype, device=vv.device)) and
replace vv.device with vv.place so it becomes device=vv.place; also scan nearby
uses of vv.device and update them similarly if present.
In `@deepmd/pd/model/task/fitting.py`:
- Around line 128-141: Replace the incorrect paddle.tensor(...) calls used when
copying into base_class.aparam_avg and base_class.aparam_inv_std with
paddle.to_tensor(...) and pass the appropriate place argument (e.g., env.PLACE)
instead of device, preserving the dtype from base_class.aparam_avg.dtype and
base_class.aparam_inv_std.dtype and using the original aparam_avg and
aparam_inv_std arrays; update the two sites around base_class.aparam_avg.copy_
and base_class.aparam_inv_std.copy_ to create tensors via
paddle.to_tensor(aparam_avg, dtype=..., place=env.PLACE) and
paddle.to_tensor(aparam_inv_std, dtype=..., place=env.PLACE).
- Around line 536-541: Replace the incorrect paddle.tensor(...) call when
creating default_fparam_tensor with paddle.to_tensor(..., place=...) and pass a
Paddle Place for the device; specifically, in the self.register_buffer call that
creates "default_fparam_tensor" (and where self.default_fparam is converted via
np.array), change paddle.tensor(np.array(self.default_fparam), dtype=self.prec,
device=device) to paddle.to_tensor(np.array(self.default_fparam),
dtype=self.prec, place=device) (or construct a
paddle.CPUPlace()/paddle.CUDAPlace(...) into `device` beforehand if `device` is
not already a Place object) so the tensor uses the correct API and place
parameter.
- Around line 885-889: The paddle.zeros call that creates outs in
deepmd/pd/model/task/fitting.py uses an invalid device= keyword; replace it to
use the paddle API's place= parameter (matching the fix in transform_output.py).
Update the paddle.zeros call that constructs outs (nf, nloc, net_dim_out) to
pass place=descriptor.place (or otherwise construct a paddle.CPUPlace/CUDAPlace
from descriptor if descriptor.place is not already available) so the tensor is
allocated on the correct Paddle place. Ensure the symbol outs and the
paddle.zeros call are the only changes and keep dtype and shape the same.
- Around line 99-112: The code wrongly calls paddle.tensor(...) which doesn't
exist; change both uses inside base_class.fparam_avg.copy_(...) and
base_class.fparam_inv_std.copy_(...) to paddle.to_tensor(...), passing dtype=...
and the Paddle "place" instead of "device" (e.g., place=env.DEVICE) so the
tensors are created with the correct dtype and device/place before calling
copy_. Ensure you preserve the same source arrays fparam_avg and fparam_inv_std
and the target dtypes base_class.fparam_avg.dtype and
base_class.fparam_inv_std.dtype when constructing the paddle.to_tensor calls.
---
Outside diff comments:
In `@deepmd/pd/model/descriptor/dpa2.py`:
- Around line 937-939: The enable_compression method currently always raises
ValueError("Compression is already enabled.") without checking the actual
compression flag; update the method (enable_compression in class DescrptDPA2) to
first inspect self.compress and only raise that ValueError if self.compress is
already True, otherwise proceed to enable compression (mirror the conditional
behavior in DescrptDPA1) so the error message accurately reflects state.
---
Nitpick comments:
In `@deepmd/pd/model/atomic_model/dp_atomic_model.py`:
- Line 374: The tuple returned by self.descriptor is unpacked into descriptor,
rot_mat, g2, h2, sw but sw is never used; rename the variable to _sw (or prefix
it with an underscore) in the unpacking site inside the method using
self.descriptor so static analyzers treat it as intentionally unused and avoid
warnings.
In `@deepmd/pd/model/descriptor/dpa2.py`:
- Around line 852-860: Standardize Paddle tensor dtype casting in this module:
replace the mixed usages so all casts use the same method (either
.astype(dtype=env.GLOBAL_PD_FLOAT_PRECISION) or
.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION)) for consistency; update the return
tuple (g1, rot_mat, g2, h2, sw) to use the chosen form to match the earlier
.astype usage on the other tensors in this file (keep the None checks
unchanged).
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 297944cc-816a-418b-ab6c-a69cf83bb836
📒 Files selected for processing (11)
deepmd/pd/model/atomic_model/base_atomic_model.pydeepmd/pd/model/atomic_model/dp_atomic_model.pydeepmd/pd/model/descriptor/dpa1.pydeepmd/pd/model/descriptor/dpa2.pydeepmd/pd/model/descriptor/repflows.pydeepmd/pd/model/model/make_model.pydeepmd/pd/model/model/transform_output.pydeepmd/pd/model/task/fitting.pydeepmd/pt/model/descriptor/repflows.pydeepmd/pt/model/task/ener.pydeepmd/pt/model/task/fitting.py
💤 Files with no reviewable changes (2)
- deepmd/pt/model/descriptor/repflows.py
- deepmd/pd/model/descriptor/repflows.py
✅ Files skipped from review due to trivial changes (1)
- deepmd/pt/model/task/ener.py
Comment on lines
+301
to
+307
| self.tebd_compress = False | ||
| if type_map is not None: | ||
| self.register_buffer( | ||
| "buffer_type_map", | ||
| paddle.to_tensor([ord(c) for c in " ".join(type_map)]), | ||
| ) | ||
| self.compress = False | ||
| self.geo_compress = False |
Contributor
There was a problem hiding this comment.
Missing compression flag check before raising error.
The tebd_compress and geo_compress flags are initialized correctly, but enable_compression() unconditionally raises "Compression is already enabled" without checking these flags. The PyTorch reference implementation (deepmd/pt/model/descriptor/dpa1.py:598-599) validates the flags before raising this error.
🐛 Proposed fix for enable_compression
def enable_compression(
self,
min_nbor_dist: float,
table_extrapolate: float = 5,
table_stride_1: float = 0.01,
table_stride_2: float = 0.1,
check_frequency: int = -1,
) -> None:
# do some checks before the mocel compression process
+ if self.tebd_compress or self.geo_compress:
+ raise ValueError("Compression is already enabled.")
- raise ValueError("Compression is already enabled.")
+ raise NotImplementedError("Model compression is not yet supported.")Also applies to: 619-620
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/descriptor/dpa1.py` around lines 301 - 307, The
enable_compression method currently raises "Compression is already enabled"
unconditionally; modify enable_compression to first check the instance flags
self.tebd_compress and self.geo_compress (and any other compression flags used
in this class) and only raise that error if both relevant flags are already
True; otherwise set the appropriate flags and proceed to enable compression.
Also apply the same conditional-check pattern to the other symmetric location
that raises the same error (the second occurrence referring to the same
enable_compression logic) so both places validate
self.tebd_compress/self.geo_compress before raising.
Comment on lines
466
to
470
| * paddle.ones( | ||
| [n_nf, n_nloc, nnei - n_nnei], | ||
| dtype=nlist.dtype, | ||
| device=nlist.device, | ||
| ), |
Contributor
There was a problem hiding this comment.
paddle.zeros() uses place not device parameter.
Consistent with issues found in other files.
🐛 Proposed fix
nlist = paddle.concat(
[
nlist,
-1
* paddle.zeros(
[n_nf, n_nloc, nnei - n_nnei],
dtype=nlist.dtype,
- device=nlist.device,
- ),
+ ).to(device=nlist.place),
],
axis=-1,
)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/model/make_model.py` around lines 466 - 470, In the paddle
tensor creation call that currently uses device=nlist.device (the
paddle.ones([...], dtype=nlist.dtype, device=nlist.device) expression involving
n_nf, n_nloc, nnei, n_nnei and nlist), replace the incorrect device keyword with
paddle's place keyword (e.g., place=nlist.place) so the call becomes
paddle.ones(..., dtype=nlist.dtype, place=nlist.place); also audit nearby
paddle.zeros/paddle.ones calls in the same scope and switch any device=...
usages to place=... while keeping dtype and shape unchanged.
Comment on lines
+239
to
+241
| force = paddle.zeros( | ||
| vldims + derv_r_ext_dims, dtype=vv.dtype, device=vv.device | ||
| ) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🌐 Web query:
What is the correct API for paddle.zeros to specify device placement in PaddlePaddle?
💡 Result:
paddle.zeros uses the keyword-only argument device (not place) to control device placement:
x = paddle.zeros([2, 3], device="gpu:0") # or "cpu"
# or
x = paddle.zeros([2, 3], device=paddle.CUDAPlace(0))If you leave device=None, it creates the tensor on Paddle’s current global device, which you can set via:
paddle.device.set_device("gpu:0")
x = paddle.zeros([2, 3]) # now on gpu:0Sources: Paddle paddle.zeros API (shows device: PlaceLike) [1], Paddle paddle.device.set_device API (shows "gpu:0" form) [2].
🏁 Script executed:
#!/bin/bash
# First, check if the file exists and read the context around lines 239-241
if [ -f "deepmd/pd/model/model/transform_output.py" ]; then
echo "=== File context (lines 230-250) ==="
sed -n '230,250p' deepmd/pd/model/model/transform_output.py | cat -n
else
echo "File not found at: deepmd/pd/model/model/transform_output.py"
# Try to find it
find . -name "transform_output.py" -type f
fiRepository: deepmodeling/deepmd-kit
Length of output: 1123
🏁 Script executed:
#!/bin/bash
# Search for paddle.zeros usage in the codebase to see common patterns
echo "=== Searching for paddle.zeros usage patterns ==="
rg "paddle\.zeros\(" -A 2 -B 1 | head -100Repository: deepmodeling/deepmd-kit
Length of output: 6312
🏁 Script executed:
#!/bin/bash
# Search for vv.device and vv.place usage patterns
echo "=== Searching for .device and .place patterns ==="
rg "\.device|\.place" deepmd/pd/model/model/transform_output.py -B 2 -A 1 | head -60Repository: deepmodeling/deepmd-kit
Length of output: 426
🏁 Script executed:
#!/bin/bash
# Read more context around line 239-241 to understand what vv is
sed -n '200,260p' deepmd/pd/model/model/transform_output.py | cat -nRepository: deepmodeling/deepmd-kit
Length of output: 2915
🏁 Script executed:
#!/bin/bash
# Search for how vv is defined in this file
rg "^\s*vv\s*=" deepmd/pd/model/model/transform_output.py -B 5 -A 2Repository: deepmodeling/deepmd-kit
Length of output: 275
🏁 Script executed:
#!/bin/bash
# Check the imports and class definition
head -50 deepmd/pd/model/model/transform_output.py | cat -nRepository: deepmodeling/deepmd-kit
Length of output: 1769
🏁 Script executed:
#!/bin/bash
# Search for tensor.place patterns in the codebase
echo "=== Searching for .place usage on tensors ==="
rg "\.place\b" deepmd/pd/ -B 1 -A 1 | head -80Repository: deepmodeling/deepmd-kit
Length of output: 5055
🏁 Script executed:
#!/bin/bash
# Search for .device patterns on tensors to see if it's used elsewhere
echo "=== Searching for .device usage on tensors ==="
rg "(\w+)\.device\b" deepmd/pd/ -o | sort | uniq -c | sort -rn | head -20Repository: deepmodeling/deepmd-kit
Length of output: 442
paddle.Tensor uses place not device attribute.
The Paddle API for paddle.zeros correctly accepts a device parameter, but paddle.Tensor objects have a .place attribute, not .device. The current code will raise an AttributeError.
🐛 Proposed fix
force = paddle.zeros(
- vldims + derv_r_ext_dims, dtype=vv.dtype, device=vv.device
+ vldims + derv_r_ext_dims, dtype=vv.dtype, device=vv.place🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/model/transform_output.py` around lines 239 - 241, The
paddle.zeros call is passing vv.device which doesn't exist on paddle.Tensor;
change the device argument to use vv.place instead. Locate the paddle.zeros
invocation that builds force (the line creating force = paddle.zeros(vldims +
derv_r_ext_dims, dtype=vv.dtype, device=vv.device)) and replace vv.device with
vv.place so it becomes device=vv.place; also scan nearby uses of vv.device and
update them similarly if present.
Comment on lines
+258
to
+260
| virial = paddle.zeros( | ||
| vldims + derv_c_ext_dims, dtype=vv.dtype, device=vv.device | ||
| ) |
Contributor
There was a problem hiding this comment.
Same issue: paddle.zeros uses place not device.
Consistent with the earlier issue at line 240.
🐛 Proposed fix
virial = paddle.zeros(
- vldims + derv_c_ext_dims, dtype=vv.dtype, device=vv.device
+ vldims + derv_c_ext_dims, dtype=vv.dtype
+ ).to(device=vv.place)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/model/transform_output.py` around lines 258 - 260, The
paddle.zeros call creating virial uses the wrong keyword argument 'device';
change it to use 'place' (matching Paddle API and the earlier fix at line ~240)
so the call uses place=vv.place (or equivalent) instead of device=vv.device;
update the virial creation in transform_output (the paddle.zeros(...) that
references vldims, derv_c_ext_dims and vv) to use place=vv.place.
Comment on lines
+99
to
+112
| base_class.fparam_avg.copy_( | ||
| paddle.tensor( | ||
| fparam_avg, | ||
| device=env.DEVICE, | ||
| dtype=base_class.fparam_avg.dtype, | ||
| ) | ||
| ) | ||
| base_class.fparam_inv_std.copy_( | ||
| paddle.tensor( | ||
| fparam_inv_std, | ||
| device=env.DEVICE, | ||
| dtype=base_class.fparam_inv_std.dtype, | ||
| ) | ||
| ) |
Contributor
There was a problem hiding this comment.
paddle.tensor() is not a valid Paddle API - use paddle.to_tensor().
Paddle does not have a paddle.tensor() function. The correct API is paddle.to_tensor(). Additionally, paddle.to_tensor() uses place parameter, not device.
🐛 Proposed fix for fparam buffer update
base_class.fparam_avg.copy_(
- paddle.tensor(
+ paddle.to_tensor(
fparam_avg,
- device=env.DEVICE,
+ place=env.DEVICE,
dtype=base_class.fparam_avg.dtype,
)
)
base_class.fparam_inv_std.copy_(
- paddle.tensor(
+ paddle.to_tensor(
fparam_inv_std,
- device=env.DEVICE,
+ place=env.DEVICE,
dtype=base_class.fparam_inv_std.dtype,
)
)Does PaddlePaddle have paddle.tensor() function or is it paddle.to_tensor()?
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/task/fitting.py` around lines 99 - 112, The code wrongly
calls paddle.tensor(...) which doesn't exist; change both uses inside
base_class.fparam_avg.copy_(...) and base_class.fparam_inv_std.copy_(...) to
paddle.to_tensor(...), passing dtype=... and the Paddle "place" instead of
"device" (e.g., place=env.DEVICE) so the tensors are created with the correct
dtype and device/place before calling copy_. Ensure you preserve the same source
arrays fparam_avg and fparam_inv_std and the target dtypes
base_class.fparam_avg.dtype and base_class.fparam_inv_std.dtype when
constructing the paddle.to_tensor calls.
Comment on lines
+128
to
+141
| base_class.aparam_avg.copy_( | ||
| paddle.tensor( | ||
| aparam_avg, | ||
| device=env.DEVICE, | ||
| dtype=base_class.aparam_avg.dtype, | ||
| ) | ||
| ) | ||
| base_class.aparam_inv_std.copy_( | ||
| paddle.tensor( | ||
| aparam_inv_std, | ||
| device=env.DEVICE, | ||
| dtype=base_class.aparam_inv_std.dtype, | ||
| ) | ||
| ) |
Contributor
There was a problem hiding this comment.
Same issue: paddle.tensor() should be paddle.to_tensor() with place parameter.
🐛 Proposed fix for aparam buffer update
base_class.aparam_avg.copy_(
- paddle.tensor(
+ paddle.to_tensor(
aparam_avg,
- device=env.DEVICE,
+ place=env.DEVICE,
dtype=base_class.aparam_avg.dtype,
)
)
base_class.aparam_inv_std.copy_(
- paddle.tensor(
+ paddle.to_tensor(
aparam_inv_std,
- device=env.DEVICE,
+ place=env.DEVICE,
dtype=base_class.aparam_inv_std.dtype,
)
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| base_class.aparam_avg.copy_( | |
| paddle.tensor( | |
| aparam_avg, | |
| device=env.DEVICE, | |
| dtype=base_class.aparam_avg.dtype, | |
| ) | |
| ) | |
| base_class.aparam_inv_std.copy_( | |
| paddle.tensor( | |
| aparam_inv_std, | |
| device=env.DEVICE, | |
| dtype=base_class.aparam_inv_std.dtype, | |
| ) | |
| ) | |
| base_class.aparam_avg.copy_( | |
| paddle.to_tensor( | |
| aparam_avg, | |
| place=env.DEVICE, | |
| dtype=base_class.aparam_avg.dtype, | |
| ) | |
| ) | |
| base_class.aparam_inv_std.copy_( | |
| paddle.to_tensor( | |
| aparam_inv_std, | |
| place=env.DEVICE, | |
| dtype=base_class.aparam_inv_std.dtype, | |
| ) | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/task/fitting.py` around lines 128 - 141, Replace the
incorrect paddle.tensor(...) calls used when copying into base_class.aparam_avg
and base_class.aparam_inv_std with paddle.to_tensor(...) and pass the
appropriate place argument (e.g., env.PLACE) instead of device, preserving the
dtype from base_class.aparam_avg.dtype and base_class.aparam_inv_std.dtype and
using the original aparam_avg and aparam_inv_std arrays; update the two sites
around base_class.aparam_avg.copy_ and base_class.aparam_inv_std.copy_ to create
tensors via paddle.to_tensor(aparam_avg, dtype=..., place=env.PLACE) and
paddle.to_tensor(aparam_inv_std, dtype=..., place=env.PLACE).
Comment on lines
+536
to
+541
| self.register_buffer( | ||
| "default_fparam_tensor", | ||
| paddle.tensor( | ||
| np.array(self.default_fparam), dtype=self.prec, device=device | ||
| ), | ||
| ) |
Contributor
There was a problem hiding this comment.
Same issue: paddle.tensor() should be paddle.to_tensor() with place parameter.
🐛 Proposed fix for default_fparam_tensor initialization
self.register_buffer(
"default_fparam_tensor",
- paddle.tensor(
- np.array(self.default_fparam), dtype=self.prec, device=device
+ paddle.to_tensor(
+ np.array(self.default_fparam), dtype=self.prec, place=device
),
)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/task/fitting.py` around lines 536 - 541, Replace the
incorrect paddle.tensor(...) call when creating default_fparam_tensor with
paddle.to_tensor(..., place=...) and pass a Paddle Place for the device;
specifically, in the self.register_buffer call that creates
"default_fparam_tensor" (and where self.default_fparam is converted via
np.array), change paddle.tensor(np.array(self.default_fparam), dtype=self.prec,
device=device) to paddle.to_tensor(np.array(self.default_fparam),
dtype=self.prec, place=device) (or construct a
paddle.CPUPlace()/paddle.CUDAPlace(...) into `device` beforehand if `device` is
not already a Place object) so the tensor uses the correct API and place
parameter.
Comment on lines
885
to
889
| outs = paddle.zeros( | ||
| (nf, nloc, net_dim_out), | ||
| dtype=env.GLOBAL_PD_FLOAT_PRECISION, | ||
| device=descriptor.device, | ||
| ) |
Contributor
There was a problem hiding this comment.
paddle.zeros() uses place not device parameter.
This is consistent with the issue found in transform_output.py.
🐛 Proposed fix
outs = paddle.zeros(
(nf, nloc, net_dim_out),
dtype=env.GLOBAL_PD_FLOAT_PRECISION,
- device=descriptor.device,
- )
+ ).to(device=descriptor.place)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@deepmd/pd/model/task/fitting.py` around lines 885 - 889, The paddle.zeros
call that creates outs in deepmd/pd/model/task/fitting.py uses an invalid
device= keyword; replace it to use the paddle API's place= parameter (matching
the fix in transform_output.py). Update the paddle.zeros call that constructs
outs (nf, nloc, net_dim_out) to pass place=descriptor.place (or otherwise
construct a paddle.CPUPlace/CUDAPlace from descriptor if descriptor.place is not
already available) so the tensor is allocated on the correct Paddle place.
Ensure the symbol outs and the paddle.zeros call are the only changes and keep
dtype and shape the same.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (3)
deepmd/pd/model/task/ener.py (1)
49-49: Type hint could be more specific for consistency.The parent class
GeneralFittinguseslist[float] | Nonefor thedefault_fparamparameter type annotation. Consider using the same specific type for consistency across the inheritance hierarchy.Suggested improvement
- default_fparam: list | None = None, + default_fparam: list[float] | None = None,🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/model/task/ener.py` at line 49, The constructor parameter default_fparam in the Ener class is typed as a generic list | None which is less specific than the parent GeneralFitting signature; change the annotation to match the parent by using list[float] | None for default_fparam in Ener's __init__ (or wherever default_fparam is declared) to ensure consistent type hints across the inheritance hierarchy and avoid type-checker mismatches with GeneralFitting.deepmd/pd/model/descriptor/se_a.py (1)
292-299: Thefparamparameter is added but unused in this method.The
fparamparameter is accepted for interface consistency with other descriptors (e.g., DPA1/DPA2/DPA3) but is never referenced in the method body. Consider adding adel fparamstatement (similar to howmappingis handled elsewhere) or a comment to indicate this is intentional, which would also silence potential linter warnings about unused parameters.🔧 Suggested change to explicitly mark fparam as unused
) -> tuple[ paddle.Tensor, paddle.Tensor | None, paddle.Tensor | None, paddle.Tensor | None, paddle.Tensor | None, ]: """Compute the descriptor. ... """ # cast the input to internal precsion + del fparam # unused in this descriptor coord_ext = coord_ext.to(dtype=self.prec)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/model/descriptor/se_a.py` around lines 292 - 299, The new fparam parameter in the descriptor method signature is never used; to silence linters and make intent explicit, mark it as intentionally unused by adding a del fparam (or a clear inline comment like "# unused, kept for interface compatibility with DPA descriptors") near the top of the method body (similar to how mapping is handled); update the function in se_a.py where the signature includes fparam to include this single-line change so the parameter is referenced and the warning is avoided.deepmd/pd/model/descriptor/se_t_tebd.py (1)
440-447: Thefparamparameter is added but unused.Similar to the observation in
se_a.py, thefparamparameter is accepted for interface consistency but not used in the method body. Consider explicitly marking it as unused.🔧 Suggested change
""" # cast the input to internal precsion extended_coord = extended_coord.to(dtype=self.prec) - del mapping + del mapping, fparam # unused in this descriptor nframes, nloc, nnei = nlist.shape🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/model/descriptor/se_t_tebd.py` around lines 440 - 447, The fparam parameter in the method that declares "fparam: paddle.Tensor | None = None" is accepted but never used; mark it explicitly unused to avoid lint warnings and clarify intent by referencing it in the body (e.g., assign to a dummy variable or delete it), for example add a single line like "_ = fparam" or "del fparam" at the start of the method that returns the tuple of paddle.Tensor types so the parameter remains part of the interface but is visibly unused.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/pd/model/descriptor/se_a.py`:
- Around line 292-299: The new fparam parameter in the descriptor method
signature is never used; to silence linters and make intent explicit, mark it as
intentionally unused by adding a del fparam (or a clear inline comment like "#
unused, kept for interface compatibility with DPA descriptors") near the top of
the method body (similar to how mapping is handled); update the function in
se_a.py where the signature includes fparam to include this single-line change
so the parameter is referenced and the warning is avoided.
In `@deepmd/pd/model/descriptor/se_t_tebd.py`:
- Around line 440-447: The fparam parameter in the method that declares "fparam:
paddle.Tensor | None = None" is accepted but never used; mark it explicitly
unused to avoid lint warnings and clarify intent by referencing it in the body
(e.g., assign to a dummy variable or delete it), for example add a single line
like "_ = fparam" or "del fparam" at the start of the method that returns the
tuple of paddle.Tensor types so the parameter remains part of the interface but
is visibly unused.
In `@deepmd/pd/model/task/ener.py`:
- Line 49: The constructor parameter default_fparam in the Ener class is typed
as a generic list | None which is less specific than the parent GeneralFitting
signature; change the annotation to match the parent by using list[float] | None
for default_fparam in Ener's __init__ (or wherever default_fparam is declared)
to ensure consistent type hints across the inheritance hierarchy and avoid
type-checker mismatches with GeneralFitting.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: a1801a5a-f902-4db4-bf62-3ce87b441fc0
📒 Files selected for processing (3)
deepmd/pd/model/descriptor/se_a.pydeepmd/pd/model/descriptor/se_t_tebd.pydeepmd/pd/model/task/ener.py
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pd/model/descriptor/se_a.py (1)
292-331: Document thefparamparameter and clarify it's unused.The
fparamparameter is added for interface consistency with other descriptors (like DPA3), but it's not passed toself.sea.forward()or used in the implementation. This is intentional since SeA doesn't support charge/spin embedding. Consider documenting this in the docstring to prevent confusion.📝 Suggested docstring update
"""Compute the descriptor. Parameters ---------- coord_ext The extended coordinates of atoms. shape: nf x (nallx3) atype_ext The extended aotm types. shape: nf x nall nlist The neighbor list. shape: nf x nloc x nnei mapping The index mapping, not required by this descriptor. comm_dict The data needed for communication for parallel inference. + fparam + Frame parameters. Not used by this descriptor but accepted for + interface consistency with other descriptor types. Returns🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pd/model/descriptor/se_a.py` around lines 292 - 331, The docstring lacks documentation for the fparam parameter and does not state it is intentionally unused; update the descriptor docstring to add an entry for fparam explaining it exists for interface consistency with other descriptors (e.g., DPA3), that SeA does not support charge/spin embeddings, and that fparam is deliberately not passed to self.sea.forward() or used elsewhere in se_a (so callers can ignore it). Mention the parameter type (paddle.Tensor | None) and that it is accepted but unused to prevent confusion.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/pd/model/descriptor/se_a.py`:
- Around line 292-331: The docstring lacks documentation for the fparam
parameter and does not state it is intentionally unused; update the descriptor
docstring to add an entry for fparam explaining it exists for interface
consistency with other descriptors (e.g., DPA3), that SeA does not support
charge/spin embeddings, and that fparam is deliberately not passed to
self.sea.forward() or used elsewhere in se_a (so callers can ignore it). Mention
the parameter type (paddle.Tensor | None) and that it is accepted but unused to
prevent confusion.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 02b1bd91-2fae-4163-86c3-d30da3acac5c
📒 Files selected for processing (2)
deepmd/pd/model/descriptor/dpa1.pydeepmd/pd/model/descriptor/se_a.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5319 +/- ##
==========================================
- Coverage 82.30% 82.20% -0.10%
==========================================
Files 775 775
Lines 77628 77820 +192
Branches 3675 3675
==========================================
+ Hits 63888 63975 +87
- Misses 12568 12670 +102
- Partials 1172 1175 +3 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces support for charge and spin embeddings in the atomic descriptor model (
DescrptDPA3) and integrates this feature into the atomic model (DPAtomicModel). The changes improve the model's ability to handle additional atomic properties and update serialization, deserialization, and testing to accommodate these enhancements. The most important changes are grouped below by theme.Charge and Spin Embedding Support
add_chg_spin_ebdparameter toDescrptDPA3, enabling charge and spin embedding layers (chg_embedding,spin_embedding,mix_cs_mlp) and activation function. The forward method now incorporates charge and spin embeddings when enabled. (F5ea4df3L121R121, F5ea4df3L176R176, F5ea4df3L199R199, F5ea4df3L591R591)DescrptDPA3to handle charge and spin embedding layers, ensuring proper model checkpointing and restoration. (F5ea4df3L464R464, F5ea4df3L498R498)Integration with Atomic Model
add_chg_spin_ebdfrom descriptor toDPAtomicModel, initializing relevant layers and updating the constructor. (F41ff152L76R76)Testing Enhancements
test_consistencyintest_dpa3.pyto include charge and spin embedding scenarios, ensuring correctness across serialization, deserialization, and model inference. (source/tests/pd/model/test_dpa3.pyR58, F2294b18L66R66, F2294b18L104R104, F2294b18L144R144)Miscellaneous Improvements
DPAtomicModeldocstring for clarity. (F41ff152L30R30)DescrptDPA3.forwardto support optional outputs and maintain precision. (F5ea4df3L619R619)Summary by CodeRabbit
New Features
Bug Fixes
Tests
Chores